diff --git a/.github/workflows/audit.yml b/.github/workflows/audit.yml index 494e469a65ac9..e3071c1fc8ce7 100644 --- a/.github/workflows/audit.yml +++ b/.github/workflows/audit.yml @@ -45,7 +45,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-audit - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-audit - name: Run audit check diff --git a/.github/workflows/breaking_changes_detector.yml b/.github/workflows/breaking_changes_detector.yml index 9e23fef8a6c8a..a44a50a53f1e8 100644 --- a/.github/workflows/breaking_changes_detector.yml +++ b/.github/workflows/breaking_changes_detector.yml @@ -89,7 +89,7 @@ jobs: - name: Install cargo-semver-checks if: steps.changed_crates.outputs.packages != '' - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-semver-checks diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index b9f90dbd8eb14..1a7d075d9e886 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -45,11 +45,11 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4 + uses: github/codeql-action/init@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@db488ddef3bf6cb639b32c2e9a7c0a7ea8271d28 # v4 + uses: github/codeql-action/analyze@cdf488f595d80d6e07e03d4674febd5ab45fa938 # v4 with: category: "/language:actions" diff --git a/.github/workflows/dependencies.yml b/.github/workflows/dependencies.yml index d5acaa6338db7..726d0783eee39 100644 --- a/.github/workflows/dependencies.yml +++ b/.github/workflows/dependencies.yml @@ -63,7 +63,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-machete - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-machete@0.9 - name: Detect unused dependencies diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index fbbb1e2a9d691..83ebb6e78b678 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -38,7 +38,7 @@ jobs: steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install HawkEye - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: hawkeye@7.0.0 - name: Run license header check @@ -66,7 +66,7 @@ jobs: source ci/scripts/utils/tool_versions.sh echo "LYCHEE_VERSION=${LYCHEE_VERSION}" >> "$GITHUB_ENV" - name: Install lychee - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: lychee@${{ env.LYCHEE_VERSION }} - name: Run markdown link check @@ -91,7 +91,7 @@ jobs: # it may cause checks to fail more often. # We can upgrade it manually once a while. - name: Install typos - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: typos@1.37.0 - name: Run typos check diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index c52ed5681c64f..4eda5ac4590a2 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -53,7 +53,7 @@ jobs: sudo apt-get update sudo apt-get install -y graphviz - name: Install cargo-depgraph - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-depgraph@1.6 diff --git a/.github/workflows/docs_pr.yaml b/.github/workflows/docs_pr.yaml index 2cbb64274a526..097cfe130b5b0 100644 --- a/.github/workflows/docs_pr.yaml +++ b/.github/workflows/docs_pr.yaml @@ -59,7 +59,7 @@ jobs: sudo apt-get update sudo apt-get install -y graphviz - name: Install cargo-depgraph - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-depgraph@1.6 - name: Build docs html and check for warnings diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 9de0c80e7a1f0..7925fdd817f0c 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -346,7 +346,7 @@ jobs: - name: Install llvm-tools-preview run: rustup component add llvm-tools-preview - name: Install cargo-llvm-cov - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-llvm-cov - name: Rust Dependency Cache @@ -491,7 +491,7 @@ jobs: sudo apt-get update -qq sudo apt-get install -y -qq clang - name: Setup wasm-pack - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: wasm-pack - name: Run tests with headless mode @@ -717,7 +717,7 @@ jobs: with: rust-version: stable - name: Install taplo - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: taplo-cli@0.9 # if you encounter an error, try running 'taplo format' to fix the formatting automatically. @@ -802,7 +802,7 @@ jobs: - name: Setup Rust toolchain uses: ./.github/actions/setup-builder - name: Install cargo-msrv - uses: taiki-e/install-action@ba47c86ac325773530516bb756137ac718732518 # v2.86.5 + uses: taiki-e/install-action@742a3317eac7bd62f91cd888b4eead5e784ba833 # v2.87.1 with: tool: cargo-msrv diff --git a/Cargo.lock b/Cargo.lock index 7a68fcb5e63d7..293ac45c51e38 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -226,7 +226,7 @@ dependencies = [ "bzip2", "crc", "flate2", - "indexmap 2.14.0", + "indexmap 2.14.1", "liblzma", "rand 0.9.4", "serde", @@ -356,7 +356,7 @@ dependencies = [ "arrow-select", "chrono", "half", - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "lexical-core", "memchr", @@ -511,7 +511,7 @@ checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -1023,11 +1023,11 @@ checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" [[package]] name = "blake2" -version = "0.10.6" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +checksum = "5b5d4d889834ee8ecfc0f8426ad30faf7cdcb10f741a8e6d7224d95325479f6f" dependencies = [ - "digest 0.10.7", + "digest 0.11.2", ] [[package]] @@ -1244,9 +1244,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chacha20" -version = "0.10.0" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", "cpufeatures 0.3.0", @@ -1335,7 +1335,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -1510,9 +1510,9 @@ checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" [[package]] name = "crc32fast" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" dependencies = [ "cfg-if", ] @@ -1746,7 +1746,7 @@ dependencies = [ "futures", "glob", "half", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "liblzma", @@ -1897,7 +1897,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "hex", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "libc", @@ -2173,7 +2173,7 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "env_logger", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "recursive", @@ -2187,7 +2187,7 @@ version = "55.0.0" dependencies = [ "arrow", "datafusion-common", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "rstest", @@ -2367,7 +2367,7 @@ version = "55.0.0" dependencies = [ "datafusion-doc", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -2388,7 +2388,7 @@ dependencies = [ "datafusion-physical-expr", "datafusion-sql", "env_logger", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "log", @@ -2402,6 +2402,7 @@ name = "datafusion-physical-expr" version = "55.0.0" dependencies = [ "arrow", + "arrow-schema", "criterion", "datafusion-common", "datafusion-expr", @@ -2412,7 +2413,7 @@ dependencies = [ "datafusion-proto-models", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "parking_lot", @@ -2447,7 +2448,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-proto-models", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "itertools 0.15.0", "parking_lot", "pin-project", @@ -2505,7 +2506,7 @@ dependencies = [ "futures", "half", "hashbrown 0.17.1", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "log", @@ -2658,7 +2659,7 @@ dependencies = [ "datafusion-functions-nested", "datafusion-functions-window", "env_logger", - "indexmap 2.14.0", + "indexmap 2.14.1", "insta", "itertools 0.15.0", "log", @@ -2776,7 +2777,6 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer 0.10.4", "crypto-common 0.1.7", - "subtle", ] [[package]] @@ -3059,9 +3059,9 @@ dependencies = [ [[package]] name = "flate2" -version = "1.1.9" +version = "1.1.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" dependencies = [ "crc32fast", "miniz_oxide", @@ -3302,7 +3302,7 @@ dependencies = [ "futures-core", "futures-sink", "http 1.4.0", - "indexmap 2.14.0", + "indexmap 2.14.1", "slab", "tokio", "tokio-util", @@ -3734,9 +3734,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.14.0" +version = "2.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +checksum = "07aa2048142242915a31d35844fb311e0e53fcca590c3a0a40dcf1b841fa09eb" dependencies = [ "equivalent", "hashbrown 0.17.1", @@ -4072,9 +4072,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.33" +version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" [[package]] name = "lru-slab" @@ -4150,9 +4150,9 @@ dependencies = [ [[package]] name = "miniz_oxide" -version = "0.8.9" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" dependencies = [ "adler2", "simd-adler32", @@ -4635,7 +4635,7 @@ checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" dependencies = [ "fixedbitset", "hashbrown 0.15.5", - "indexmap 2.14.0", + "indexmap 2.14.1", "serde", ] @@ -5605,7 +5605,7 @@ checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -5625,7 +5625,7 @@ version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "memchr", "serde", @@ -5688,7 +5688,7 @@ dependencies = [ "chrono", "hex", "indexmap 1.9.3", - "indexmap 2.14.0", + "indexmap 2.14.1", "schemars 0.9.0", "schemars 1.2.1", "serde_core", @@ -5715,7 +5715,7 @@ version = "0.9.34+deprecated" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "itoa", "ryu", "serde", @@ -6040,7 +6040,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e620ff4d5c02fd6f7752931aa74b16a26af66a63022cc1ad412c77edbe0bab47" dependencies = [ "heck", - "indexmap 2.14.0", + "indexmap 2.14.1", "pbjson 0.8.0", "pbjson-build 0.8.0", "pbjson-types", @@ -6079,9 +6079,9 @@ dependencies = [ [[package]] name = "syn" -version = "3.0.3" +version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +checksum = "e6275cddf4610d1775e6d1fe9469b2e77d0f39fd98fb7450901b821e0c53649f" dependencies = [ "proc-macro2", "quote", @@ -6204,7 +6204,7 @@ checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", - "syn 3.0.3", + "syn 3.0.4", ] [[package]] @@ -6387,7 +6387,7 @@ version = "1.1.4+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3aace63f4bbcdfc2c965b059de67119c89c4017a70d633be6c104910f67056f5" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "serde_core", "serde_spanned", "toml_datetime", @@ -6411,7 +6411,7 @@ version = "0.25.11+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b59c4d22ed448339746c59b905d24568fcbb3ab65a500494f7b8c3e97739f2b" dependencies = [ - "indexmap 2.14.0", + "indexmap 2.14.1", "toml_datetime", "toml_parser", "winnow", @@ -6480,7 +6480,7 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", - "indexmap 2.14.0", + "indexmap 2.14.1", "pin-project-lite", "slab", "sync_wrapper", @@ -6586,9 +6586,9 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" [[package]] name = "twox-hash" -version = "2.1.3" +version = "2.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" +checksum = "5283634e518fe9e82c7b20520bb4bc209009fd16c82077c802f8111ecbb0117a" dependencies = [ "rand 0.10.1", ] @@ -6781,9 +6781,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.24.1" +version = "1.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2cefc03fd367c0c6d4305de1b312cf00248c4114f4a0418ce6a6af769e3b0bd9" +checksum = "b5772d71c9be8a8a6ac2117d949c5b224c1b72241bb611d9a3012edcf8af7812" dependencies = [ "getrandom 0.4.2", "js-sys", @@ -6989,7 +6989,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" dependencies = [ "anyhow", - "indexmap 2.14.0", + "indexmap 2.14.1", "wasm-encoder", "wasmparser", ] @@ -7015,7 +7015,7 @@ checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ "bitflags", "hashbrown 0.15.5", - "indexmap 2.14.0", + "indexmap 2.14.1", "semver", ] @@ -7392,7 +7392,7 @@ checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" dependencies = [ "anyhow", "heck", - "indexmap 2.14.0", + "indexmap 2.14.1", "prettyplease", "syn 2.0.119", "wasm-metadata", @@ -7423,7 +7423,7 @@ checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", "bitflags", - "indexmap 2.14.0", + "indexmap 2.14.1", "log", "serde", "serde_derive", @@ -7442,7 +7442,7 @@ checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" dependencies = [ "anyhow", "id-arena", - "indexmap 2.14.0", + "indexmap 2.14.1", "log", "semver", "serde", diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index df9b7f6c94f16..419fd5be3ad2b 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -139,6 +139,12 @@ h2o_big_join: h2oai benchmark with large dataset (1e9 rows) fo h2o_small_window: Extended h2oai benchmark with small dataset (1e7 rows) for window, default file format is csv h2o_medium_window: Extended h2oai benchmark with medium dataset (1e8 rows) for window, default file format is csv h2o_big_window: Extended h2oai benchmark with large dataset (1e9 rows) for window, default file format is csv +h2o_small_window_sorted: Window Top-N over a declared-sorted h2o input, small dataset (1e7 rows), default file format is csv +h2o_medium_window_sorted: Window Top-N over a declared-sorted h2o input, medium dataset (1e8 rows), default file format is csv +h2o_big_window_sorted: Window Top-N over a declared-sorted h2o input, large dataset (1e9 rows), default file format is csv +h2o_small_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, small dataset (1e7 rows), source file format is parquet +h2o_medium_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, medium dataset (1e8 rows), source file format is parquet +h2o_big_window_sorted_parquet: Window Top-N over a declared-sorted h2o input, large dataset (1e9 rows), source file format is parquet h2o_small_parquet: h2oai benchmark with small dataset (1e7 rows) for groupby, file format is parquet h2o_medium_parquet: h2oai benchmark with medium dataset (1e8 rows) for groupby, file format is parquet h2o_big_parquet: h2oai benchmark with large dataset (1e9 rows) for groupby, file format is parquet @@ -309,6 +315,26 @@ main() { h2o_big_window) data_h2o_join "BIG" "CSV" ;; + # the sorted window subgroup derives its data from the same + # source, then sorts it inside its load SQL + h2o_small_window_sorted) + data_h2o_join "SMALL" "CSV" + ;; + h2o_medium_window_sorted) + data_h2o_join "MEDIUM" "CSV" + ;; + h2o_big_window_sorted) + data_h2o_join "BIG" "CSV" + ;; + h2o_small_window_sorted_parquet) + data_h2o_join "SMALL" "PARQUET" + ;; + h2o_medium_window_sorted_parquet) + data_h2o_join "MEDIUM" "PARQUET" + ;; + h2o_big_window_sorted_parquet) + data_h2o_join "BIG" "PARQUET" + ;; h2o_small_parquet) data_h2o "SMALL" "PARQUET" ;; @@ -534,6 +560,24 @@ main() { h2o_big_window) run_h2o_window "BIG" "CSV" "window" ;; + h2o_small_window_sorted) + run_h2o_window_sorted "small" "csv" + ;; + h2o_medium_window_sorted) + run_h2o_window_sorted "medium" "csv" + ;; + h2o_big_window_sorted) + run_h2o_window_sorted "big" "csv" + ;; + h2o_small_window_sorted_parquet) + run_h2o_window_sorted "small" "parquet" + ;; + h2o_medium_window_sorted_parquet) + run_h2o_window_sorted "medium" "parquet" + ;; + h2o_big_window_sorted_parquet) + run_h2o_window_sorted "big" "parquet" + ;; h2o_small_parquet) run_h2o "SMALL" "PARQUET" ;; @@ -1266,6 +1310,38 @@ run_h2o_window() { h2o_runner "$1" "$2" "window" } +# Runs the h2o window_sorted subgroup: window Top-N over an input that declares +# the ordering the window requires. +# +# The `window` subgroup registers `x` with no declared ordering, so +# `output_ordering()` is None and any plan that depends on a declared ordering is +# unreachable from it, however the data happens to sit on disk. This subgroup's +# load SQL writes a sorted copy and registers it `WITH ORDER`, then asserts a +# per-partition top-K operator is actually in the plan so a silent fallback to +# window-plus-filter cannot masquerade as a result. +# +# The sort happens in the untimed `load` step, so it stays out of the +# measurement. Data comes from the same source as the window subgroup, so the +# data step is data_h2o_join. +# +# H2O_FILE_TYPE selects the *source* format and so only affects that untimed +# load: the measured query always reads the sorted Parquet copy, whatever the +# source was. A Parquet source still makes the load markedly cheaper than +# re-parsing a multi-GB CSV, which is why the `_parquet` entries exist; they are +# expected to produce the same measured numbers as their CSV counterparts, not +# different ones. Each entry pairs with the data step that generates its format. +run_h2o_window_sorted() { + SIZE=${1:-"small"} + FILE_TYPE=${2:-"csv"} + echo "Running h2o window_sorted benchmark (size=${SIZE}, source format=${FILE_TYPE})..." + debug_run env BENCH_NAME=h2o \ + BENCH_SUBGROUP=window_sorted \ + H2O_BENCH_SIZE="${SIZE}" \ + H2O_FILE_TYPE="${FILE_TYPE}" \ + ${QUERY:+BENCH_QUERY="${QUERY}"} \ + bash -c "$SQL_CARGO_COMMAND" +} + # Runs the external aggregation benchmark run_external_aggr() { # Use TPC-H SF1 dataset diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark new file mode 100644 index 0000000000000..236220b7595b4 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q01.benchmark @@ -0,0 +1,13 @@ +# ROW_NUMBER top-2 per partition over the declared-sorted input, swept across +# partition cardinality. This mirrors the `window` subgroup q14-q17 sweep, with +# the partition key materialized as a column so the ordering can be declared. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=01 +FN=ROW_NUMBER +PK_MOD=100 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q01_row_number_100_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark new file mode 100644 index 0000000000000..a9c5733d62e43 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q02.benchmark @@ -0,0 +1,11 @@ +# ROW_NUMBER top-2, ~1K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=02 +FN=ROW_NUMBER +PK_MOD=1000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q02_row_number_1k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark new file mode 100644 index 0000000000000..4a000af56bbf6 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q03.benchmark @@ -0,0 +1,11 @@ +# ROW_NUMBER top-2, ~10K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=03 +FN=ROW_NUMBER +PK_MOD=10000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q03_row_number_10k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark new file mode 100644 index 0000000000000..77208a20249ae --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q04.benchmark @@ -0,0 +1,12 @@ +# ROW_NUMBER top-2, ~100K partitions. Fewest rows per partition on the sweep, +# so per-partition bookkeeping is amortized over the least work. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=04 +FN=ROW_NUMBER +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q04_row_number_100k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark new file mode 100644 index 0000000000000..e06acffd9a440 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q05.benchmark @@ -0,0 +1,15 @@ +# RANK top-2 per partition, mirroring the ROW_NUMBER sweep above. RANK retains +# boundary ties, so `WHERE rnk <= 2` can keep more than 2 rows per partition and +# the operator has to track the tied set rather than a fixed row count. `v2` is +# near-unique, so these four measure RANK with the tie machinery idle; q09-q10 +# drive it. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=05 +FN=RANK +PK_MOD=100 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q05_rank_100_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark new file mode 100644 index 0000000000000..d9421902b412a --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q06.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~1K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=06 +FN=RANK +PK_MOD=1000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q06_rank_1k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark new file mode 100644 index 0000000000000..7a2a49ebce84e --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q07.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~10K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=07 +FN=RANK +PK_MOD=10000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q07_rank_10k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark new file mode 100644 index 0000000000000..409ed17bba712 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q08.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~100K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=08 +FN=RANK +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q08_rank_100k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark new file mode 100644 index 0000000000000..9728f470819b0 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q09.benchmark @@ -0,0 +1,17 @@ +# RANK top-2 with heavy ties: 10 distinct ORDER BY values, so a whole partition +# slice ties at rank 1 and every tied row survives `rnk <= 2`. The tied set, not +# the top-K bound, dominates the work. +# +# The cast is load-bearing. `v2` is Float64, so `v2 % 10` is float modulo and +# stays near-unique — it produces no ties whatsoever. Truncating to an integer +# first is what actually collapses the domain to 10 values. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=09 +FN=RANK +PK_MOD=1000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q09_rank_1k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark new file mode 100644 index 0000000000000..232e3f1b3b17d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q10.benchmark @@ -0,0 +1,11 @@ +# RANK top-2, ~10K partitions, heavy ties. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=10 +FN=RANK +PK_MOD=10000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q10_rank_10k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark new file mode 100644 index 0000000000000..d5d5d6b256927 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q11.benchmark @@ -0,0 +1,13 @@ +# DENSE_RANK top-2 per partition, mirroring the RANK sweep. DENSE_RANK keeps +# every row whose ORDER BY value is among the 2 distinct-greatest in the +# partition, so rows kept per partition is unbounded in rows-per-distinct-value. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=11 +FN=DENSE_RANK +PK_MOD=100 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q11_dense_rank_100_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark new file mode 100644 index 0000000000000..c8d277ff040d6 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q12.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~1K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=12 +FN=DENSE_RANK +PK_MOD=1000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q12_dense_rank_1k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark new file mode 100644 index 0000000000000..be7246232cc56 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q13.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~10K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=13 +FN=DENSE_RANK +PK_MOD=10000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q13_dense_rank_10k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark new file mode 100644 index 0000000000000..9c773a93ea206 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q14.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~100K partitions. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=14 +FN=DENSE_RANK +PK_MOD=100000 +OB_EXPR=v2 +OB_BOUND=count(DISTINCT ob) > 1000 +NAME=window_sorted_q14_dense_rank_100k_partitions \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark new file mode 100644 index 0000000000000..3a655c0b4d397 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q15.benchmark @@ -0,0 +1,13 @@ +# DENSE_RANK top-2, ~1K partitions, heavy ties. With 10 distinct ORDER BY values +# the top-2 distinct values cover roughly a fifth of each partition, so this is +# the heaviest surviving-row count on the sweep. See q09 on why the cast matters. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=15 +FN=DENSE_RANK +PK_MOD=1000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q15_dense_rank_1k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark new file mode 100644 index 0000000000000..65baaed851b26 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window_sorted/q16.benchmark @@ -0,0 +1,11 @@ +# DENSE_RANK top-2, ~10K partitions, heavy ties. + +subgroup window_sorted + +template sql_benchmarks/h2o/window_sorted.benchmark.template +QPAD=16 +FN=DENSE_RANK +PK_MOD=10000 +OB_EXPR=CAST(v2 AS BIGINT) % 10 +OB_BOUND=count(DISTINCT ob) = 10 +NAME=window_sorted_q16_dense_rank_10k_partitions_heavy_ties \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/h2o.suite b/benchmarks/sql_benchmarks/h2o/h2o.suite index 27d83285ba026..cb2314b13c6e4 100644 --- a/benchmarks/sql_benchmarks/h2o/h2o.suite +++ b/benchmarks/sql_benchmarks/h2o/h2o.suite @@ -31,3 +31,7 @@ description = "Run H2O window query 3." [[examples]] command = "cargo run --release --bin benchmark_runner -- h2o --subgroup join --size medium -f parquet" description = "Run the H2O join queries with the medium Parquet dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o --subgroup window_sorted" +description = "Run the window Top-N sweep over an input that declares its ordering." \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql new file mode 100644 index 0000000000000..6e8f2a67b79d1 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_sorted.sql @@ -0,0 +1,40 @@ +-- Materializes a pre-sorted copy of the h2o window table that *declares* its +-- ordering, for the `window_sorted` subgroup. +-- +-- Why this exists: the `window` subgroup partitions by an expression +-- (`id3 % N`) and, in its heavy-ties queries, orders by another (`v2 % 10`). +-- `WITH ORDER` can only name columns, so both are materialized here as `pk` +-- and `ob`. The queries then partition and order by plain columns, and the +-- declared ordering can satisfy the window's requirement. +-- +-- Parameters, supplied by window_sorted.benchmark.template so each benchmark +-- writes only the shape it measures: +-- PK_MOD partition count (100, 1000, 10000, 100000) +-- OB_EXPR ORDER BY value expression (`v2` distinct, integer-cast for ties) +-- +-- Note `v2` is Float64 here, so a plain `v2 % 10` is float modulo and stays +-- near-unique. Producing genuine ties needs an integer cast first; the template +-- asserts ob's resulting cardinality so a query cannot get this wrong silently. +-- +-- The ORDER BY is load-bearing: the file has to be written in (pk, ob DESC) +-- order for the WITH ORDER declaration below to be true. A declaration that +-- does not match the file is not a planning error, it silently produces wrong +-- results, so this ORDER BY and that WITH ORDER must be edited together. +-- +-- `load` runs before `init` and is not timed, so this sort never enters the +-- measurement. +DROP TABLE IF EXISTS x_sorted; + +COPY ( + SELECT id3 % ${PK_MOD:-1000} AS pk, ${OB_EXPR:-v2} AS ob + FROM x + WHERE v2 IS NOT NULL + ORDER BY pk ASC, ob DESC +) +TO 'sql_benchmarks/h2o/scratch/window_sorted.parquet' +STORED AS PARQUET; + +CREATE EXTERNAL TABLE x_sorted +STORED AS PARQUET +WITH ORDER (pk ASC, ob DESC) +LOCATION 'sql_benchmarks/h2o/scratch/window_sorted.parquet'; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql b/benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql new file mode 100644 index 0000000000000..9601519d48f7c --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/window_sorted_cleanup.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS x_sorted; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql b/benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql new file mode 100644 index 0000000000000..7ed87226149c9 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/window_sorted_settings.sql @@ -0,0 +1,24 @@ +-- Session settings for the h2o window_sorted subgroup. `init` runs after +-- `load`, so these apply to the measured query and not to the COPY that writes +-- the sorted file. +-- +-- enable_window_topn turns on the WindowTopN rewrite, making the baseline the +-- existing per-partition top-K operator rather than an unoptimized +-- window-plus-filter plan. The flag defaults to false, so without this the +-- subgroup would measure a different plan shape than its name claims. +-- +-- information_schema is enabled so the template can assert the flags actually +-- took effect. +-- +-- prefer_existing_sort decides *how* a declared input ordering is honored once +-- some operator requires it. The heap top-K operator requires no input ordering, +-- so on its own the ordering is simply dropped: the hash repartition that +-- co-locates each partition key is inserted with preserve_order: false, and that +-- is not a defect. For an operator that does require the ordering, the planner +-- then has two ways to satisfy it — an order-preserving repartition, or a +-- SortExec that re-sorts data already sorted on disk. This setting picks the +-- former. Without it such an operator would measure a full sort and the sorted +-- input would be pointless. +set datafusion.catalog.information_schema = true; +set datafusion.optimizer.enable_window_topn = true; +set datafusion.optimizer.prefer_existing_sort = true; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template b/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template new file mode 100644 index 0000000000000..9170511492ea5 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/window_sorted.benchmark.template @@ -0,0 +1,99 @@ +# Shared template for the h2o `window_sorted` subgroup: WindowTopN over an +# input that *declares* the ordering the window requires. +# +# Parameters set by each qNN.benchmark: +# QPAD zero-padded query id +# FN window function (ROW_NUMBER, RANK, DENSE_RANK) +# PK_MOD partition count (100, 1000, 10000, 100000) +# OB_EXPR ORDER BY value expression (`v2` distinct, integer-cast for ties) +# OB_BOUND SQL predicate pinning ob's cardinality, asserted below +# NAME display name +# +# Why this subgroup exists: the `window` subgroup registers `x` with no declared +# ordering, so `output_ordering()` is None and the planner cannot know the input +# is sorted. Any operator that requires a declared ordering is therefore +# unreachable from those queries, however the data happens to be laid out on +# disk. This subgroup republishes the same sweep over a `WITH ORDER` table so +# the ordering-dependent plans are reachable. +# +# The first `load` is the existing unsorted loader, which creates `x`; the +# second reads `x` and writes the sorted `x_sorted`. Reusing the existing loader +# keeps both the `--size` and `--format` axes working with no duplication. + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +load sql_benchmarks/h2o/init/load_window_sorted.sql + +init sql_benchmarks/h2o/init/window_sorted_settings.sql + +name ${NAME} +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data, sorted by (pk, ob DESC) over ${PK_MOD} partitions + +# Guard against a silent config no-op: with the rewrite off the plan is a plain +# window plus filter, and the subgroup would measure something other than what +# its name claims. +assert I +SELECT value = 'true' FROM information_schema.df_settings WHERE name = 'datafusion.optimizer.enable_window_topn'; +---- +true + +# Same guard for the setting that decides an ordering-requiring operator gets an +# order-preserving repartition instead of a redundant SortExec. +assert I +SELECT value = 'true' FROM information_schema.df_settings WHERE name = 'datafusion.optimizer.prefer_existing_sort'; +---- +true + +# The sorted copy must hold every non-null-v2 row of the source. This catches a +# COPY that silently wrote a partial file, which would otherwise show up only as +# a suspiciously fast run. +assert I +SELECT count(*) = (SELECT count(*) FROM x WHERE v2 IS NOT NULL) FROM x_sorted; +---- +true + +# The ORDER BY column really has the cardinality this query's name claims. A +# tie-shape query whose `ob` turns out near-unique measures the no-ties case +# under a heavy-ties name, and nothing else in the run would show it: the timings +# stay plausible and only the output row count quietly collapses to K per +# partition. `v2` is Float64, so this is a live trap — `v2 % 10` is float modulo +# and yields no ties at all. +assert I +SELECT ${OB_BOUND} FROM x_sorted; +---- +true + +# The per-partition top-K operator has to be in the plan, otherwise the +# declared ordering was not picked up and this is just a window-plus-filter +# measurement wearing the subgroup's name. +# +# `PartitionedTopKExec` is deliberately the string asserted rather than a +# fully-qualified operator name: it is also a substring of the streaming +# variant's name, so this same check holds on a branch where a streaming +# operator replaces the heap one for declared-sorted input. That is the point of +# the subgroup — the comparison is this file's numbers across two branches, so +# the assertion must not be the thing that differs between them. +expect_plan PartitionedTopKExec + +# `PartitionedTopKExec` alone only proves the plan has an operator that +# distributes by the partition key; that operator requires no input ordering, +# so it appears whether or not the scan below it declares one. Assert the +# scan's own declared ordering too, so a `WITH ORDER` that silently stops +# reaching `DataSourceExec` (e.g. a rewrite that drops it, or a future change +# to `load_window_sorted.sql`) fails loudly here instead of just being absent +# from a plan that still "passes" the check above. +expect_plan output_ordering=[pk@0 ASC NULLS LAST, ob@1 DESC] + +run +-- Top-2 per partition by ${FN} over the declared-sorted input. +SELECT pk, ob FROM ( + SELECT pk, ob, + ${FN}() OVER (PARTITION BY pk ORDER BY ob DESC) AS rnk + FROM x_sorted +) sub_query WHERE rnk <= 2; + +result sql_benchmarks/h2o/results/window_sorted/${H2O_BENCH_SIZE:-small}/q${QPAD}.csv + +cleanup sql_benchmarks/h2o/init/window_sorted_cleanup.sql \ No newline at end of file diff --git a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs index 71addc6d1bcb0..202c0a71257e9 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_casts.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_casts.rs @@ -188,10 +188,11 @@ impl PhysicalExprAdapter for CustomCastsPhysicalExprAdapter { if let Some(cast) = expr.downcast_ref::() { let input_data_type = cast.expr().data_type(&self.physical_file_schema)?; - let output_data_type = cast.target_field().data_type(); + let output_field = cast.target_field(); if !cast.is_bigger_cast(&input_data_type) { return not_impl_err!( - "Unsupported CAST from {input_data_type} to {output_data_type}" + "Unsupported CAST from {input_data_type} to {}", + output_field.data_type() ); } } diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 1299e5fb7bd65..ed3dc5ea838b9 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -2614,7 +2614,7 @@ impl DataFrame { /// # async fn main() -> Result<()> { /// let id: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); /// let name: ArrayRef = Arc::new(StringArray::from(vec!["foo", "bar", "baz"])); - /// let df = DataFrame::from_columns(vec![("id", id), ("name", name)])?; + /// let df = DataFrame::from_columns([("id", id), ("name", name)])?; /// let expected = vec![ /// "+----+------+", /// "| id | name |", @@ -2628,17 +2628,16 @@ impl DataFrame { /// # Ok(()) /// # } /// ``` - pub fn from_columns(columns: Vec<(&str, ArrayRef)>) -> Result { - let fields = columns - .iter() - .map(|(name, array)| Field::new(*name, array.data_type().clone(), true)) - .collect::>(); - - let arrays = columns + pub fn from_columns<'a, I>(columns: I) -> Result + where + I: IntoIterator, + { + let (fields, arrays): (Vec<_>, Vec<_>) = columns .into_iter() - .map(|(_, array)| array) - .collect::>(); - + .map(|(name, array)| { + (Field::new(name, array.data_type().clone(), true), array) + }) + .unzip(); let schema = Arc::new(Schema::new(fields)); let batch = RecordBatch::try_new(schema, arrays)?; let ctx = SessionContext::new(); @@ -2695,7 +2694,7 @@ macro_rules! dataframe { use datafusion::prelude::DataFrame; use datafusion::common::test_util::IntoArrayRef; - let columns = vec![ + let columns = [ $( ($name, $data.into_array_ref()), )+ diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 1676a69de86d3..527c9567495e9 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -6968,7 +6968,7 @@ async fn test_dataframe_from_columns() -> Result<()> { let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); - let df = DataFrame::from_columns(vec![ + let columns = [ ("bool", bools), ("i8", i8s), ("i16", i16s), @@ -6982,10 +6982,10 @@ async fn test_dataframe_from_columns() -> Result<()> { ("f32", f32s), ("f64", f64s), ("str", strings), - ])?; + ]; - assert_eq!(df.schema().fields().len(), 13); - assert_eq!(df.clone().count().await?, 3); + let df1 = DataFrame::from_columns(columns.clone())?; + let df2 = DataFrame::from_columns(columns.to_vec())?; let expected_types = [ ("bool", DataType::Boolean), @@ -7003,14 +7003,89 @@ async fn test_dataframe_from_columns() -> Result<()> { ("str", DataType::Utf8), ]; - let schema = df.schema(); + for df in [df1, df2] { + assert_eq!(df.schema().fields().len(), expected_types.len()); + assert_eq!(df.clone().count().await?, 3); - for (name, data_type) in expected_types { - assert_eq!(schema.field_with_name(None, name)?.data_type(), &data_type); + let schema = df.schema(); + + for (name, data_type) in &expected_types { + assert_eq!(schema.field_with_name(None, name)?.data_type(), data_type); + } + + let rows = df.sort(vec![col("i32").sort(true, true)])?; + + assert_batches_eq!( + &[ + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", + "| false | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2.0 | 2.0 | bar |", + "| true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + ], + &rows.collect().await? + ); } - let rows = df.sort(vec![col("i32").sort(true, true)])?; + Ok(()) +} + +#[test] +fn test_dataframe_from_columns_empty() { + let result = DataFrame::from_columns(vec![]); + assert!(result.is_err()); + + let result = DataFrame::from_columns([]); + assert!(result.is_err()); +} + +#[tokio::test] +async fn test_dataframe_from_columns_with_iterator() -> Result<()> { + let bools: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true])); + let i8s: ArrayRef = Arc::new(Int8Array::from(vec![-1, 0, 1])); + let i16s: ArrayRef = Arc::new(Int16Array::from(vec![-1, 0, 1])); + let i32s: ArrayRef = Arc::new(Int32Array::from(vec![-1, 0, 1])); + let i64s: ArrayRef = Arc::new(Int64Array::from(vec![-1, 0, 1])); + + let u8s: ArrayRef = Arc::new(UInt8Array::from(vec![0, 1, 2])); + let u16s: ArrayRef = Arc::new(UInt16Array::from(vec![0, 1, 2])); + let u32s: ArrayRef = Arc::new(UInt32Array::from(vec![0, 1, 2])); + let u64s: ArrayRef = Arc::new(UInt64Array::from(vec![0, 1, 2])); + + let f16s: ArrayRef = Arc::new(Float16Array::from(vec![ + half::f16::from_f64(1.0), + half::f16::from_f64(2.0), + half::f16::from_f64(3.0), + ])); + let f32s: ArrayRef = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0])); + let f64s: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + + let strings: ArrayRef = + Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); + + let columns = [ + ("bool", bools), + ("i8", i8s), + ("i16", i16s), + ("i32", i32s), + ("i64", i64s), + ("u8", u8s), + ("u16", u16s), + ("u32", u32s), + ("u64", u64s), + ("f16", f16s), + ("f32", f32s), + ("f64", f64s), + ("str", strings), + ]; + let df = DataFrame::from_columns(columns.into_iter())?; + + assert_eq!(df.schema().fields().len(), 13); + assert_eq!(df.clone().count().await?, 3); + let rows = df.sort(vec![col("i32").sort(true, true)])?; assert_batches_eq!( &[ "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 15b224d200bf4..f873f956010e1 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,7 +26,10 @@ mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; -use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray}; +use arrow::array::{ + ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringArray, + StringViewArray, +}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow_schema::{DataType, Field, Schema}; @@ -125,6 +128,163 @@ async fn group_by_hash() { .await } +/// With `force_hash_collisions` every key hashes alike, so the hash +/// repartitioning sends all groups to a single final stage, whose table then +/// does not fit the limit however well memory is released. The limit is sized +/// for the real distribution across four final stages, so this test is skipped +/// under that feature. +#[cfg(not(feature = "force_hash_collisions"))] +mod count_distinct_spill { + use super::*; + use arrow::array::Int64Array; + use datafusion::assert_batches_sorted_eq; + + /// `count(distinct)` over integers under a memory limit. + /// + /// The integer distinct-count groups accumulator reports the capacity of its + /// buffers in `size()`. After an aggregate stream emits all groups, either to + /// emit partial state early or to spill, it resizes its reservation to the + /// table's reported size and expects it to have shrunk. If the accumulator + /// keeps its capacity, that resize is a grow against an exhausted pool and the + /// query fails although everything was already written out. + const COUNT_DISTINCT_ROWS: usize = 200_000; + const COUNT_DISTINCT_GROUPS: i64 = 64; + const COUNT_DISTINCT_BATCH_ROWS: usize = 8_192; + + /// Far below the distinct sets (200k values, several megabytes across the + /// partial and final tables), far above the fixed cost of the stages. With the + /// accumulator releasing its buffers the query passes from 2 MB upwards; + /// without, it fails up to 4 MB with "Decreasing allocation after spilling + /// should succeed" in the final stage or a failed emit in the partial stage. + const COUNT_DISTINCT_MEMORY_LIMIT: usize = 4 * 1024 * 1024; + + /// `g` has 64 groups, `v` is unique, so every group holds 3125 distinct values. + fn count_distinct_table() -> MemTable { + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::Int64, false), + ])); + let batches = (0..COUNT_DISTINCT_ROWS) + .step_by(COUNT_DISTINCT_BATCH_ROWS) + .map(|start| { + let rows = + start..(start + COUNT_DISTINCT_BATCH_ROWS).min(COUNT_DISTINCT_ROWS); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from_iter_values( + rows.clone().map(|row| row as i64 % COUNT_DISTINCT_GROUPS), + )), + Arc::new(Int64Array::from_iter_values( + rows.map(|row| row as i64), + )), + ], + ) + .unwrap() + }) + .collect(); + MemTable::try_new(schema, vec![batches]).unwrap() + } + + /// Four partial stages emit their state early and four hash-partitioned final + /// stages spill; every one of them must see the accumulator memory drop after + /// emitting all groups. + #[tokio::test] + async fn count_distinct_releases_memory_after_emitting_all() { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(COUNT_DISTINCT_MEMORY_LIMIT, 1.0) + .with_disk_manager_builder(DiskManagerBuilder::default()) + .build_arc() + .unwrap(); + let config = SessionConfig::new().with_target_partitions(4); + let ctx = SessionContext::new_with_config_rt(config, runtime); + ctx.register_table("t", Arc::new(count_distinct_table())) + .unwrap(); + + let batches = ctx + .sql("select count(distinct v) as d, count(*) as n from t group by g") + .await + .unwrap() + .collect() + .await + .unwrap_or_else(|error| { + panic!("query failed under the memory limit: {error}") + }); + + let per_group = (COUNT_DISTINCT_ROWS as i64 / COUNT_DISTINCT_GROUPS).to_string(); + let row = format!("| {per_group} | {per_group} |"); + let mut expected = vec!["+------+------+", "| d | n |", "+------+------+"]; + expected.extend(std::iter::repeat_n( + row.as_str(), + COUNT_DISTINCT_GROUPS as usize, + )); + expected.push("+------+------+"); + assert_batches_sorted_eq!(expected, &batches); + } +} + +/// A grouped `COUNT(DISTINCT )` gets one accumulator per group, and +/// each of those owns a hash set of the distinct values it has seen. Those +/// sets used to be created pre-allocated, which costs far more than the +/// handful of values a group typically holds, so the query's memory use +/// tracked the number of groups rather than the amount of data. +/// +/// With 4,000 groups holding 2 distinct values each, this query needed about +/// 35.5 MB of budget before the per group pre-allocation was removed and +/// about 1.9 MB after, so an 8 MB limit is a failure before the change and a +/// success after it. Spilling is disabled, so completing means the query +/// genuinely fit in the budget. +/// +/// The `avg(payload)` is load bearing, and `avg` specifically. Without a +/// second aggregate, `single_distinct_aggregation_to_group_by` rewrites the +/// distinct aggregate into a plain two stage `GROUP BY` that does not use +/// these accumulators at all. That rule tolerates a non-distinct `sum`, `min` +/// or `max` beside the distinct aggregate, because it re-aggregates its own +/// partial results over the deduplicated inner group by, and those three +/// compose with themselves. `avg` does not: averaging per group averages of +/// different sizes gives the wrong answer, so the rule can never accept it. +/// That is why ClickBench Q9 keeps its distinct aggregate. Do not replace +/// this with `count(*)`: `count` is only incidentally rejected today, and +/// proposes accepting it, +/// which would rewrite the query and leave this test passing by construction. +#[tokio::test] +async fn group_by_count_distinct_utf8() { + TestCase::new() + .with_query( + "select group_key, count(distinct value), avg(payload) from t group by group_key", + ) + .with_scenario(Scenario::GroupedDistinctStrings { + groups: 4_000, + string_view: false, + }) + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_memory_limit(8_000_000) + .with_expected_success() + .run() + .await +} + +/// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering +/// the separate view flavoured hash set. The same query over a `Utf8View` +/// column needed about 123 MB of budget before the change and about 2.5 MB +/// after, so 16 MB separates the two. +#[tokio::test] +async fn group_by_count_distinct_utf8_view() { + TestCase::new() + .with_query( + "select group_key, count(distinct value), avg(payload) from t group by group_key", + ) + .with_scenario(Scenario::GroupedDistinctStrings { + groups: 4_000, + string_view: true, + }) + .with_config(SessionConfig::new().with_target_partitions(1)) + .with_memory_limit(16_000_000) + .with_expected_success() + .run() + .await +} + #[tokio::test] async fn join_by_key_multiple_partitions() { let config = SessionConfig::new().with_target_partitions(2); @@ -982,6 +1142,14 @@ enum Scenario { /// If true, splits all input batches into 1 row each single_row_batches: bool, }, + + /// `groups` distinct integer keys paired with a short string value, for + /// grouped aggregates that build one accumulator per group. + GroupedDistinctStrings { + groups: usize, + /// If true, the value column is `Utf8View` rather than `Utf8` + string_view: bool, + }, } impl Scenario { @@ -1056,6 +1224,15 @@ impl Scenario { let table = SortedTableProvider::new(batches, sort_information); Arc::new(table) } + Self::GroupedDistinctStrings { + groups, + string_view, + } => { + let batches = grouped_distinct_string_batches(*groups, *string_view); + let table = + MemTable::try_new(batches[0].schema(), vec![batches]).unwrap(); + Arc::new(table) + } } } @@ -1079,8 +1256,62 @@ impl Scenario { // Use default rules None } + Self::GroupedDistinctStrings { .. } => { + // Disable the rules that would add a repartition, so the test + // measures the aggregate's budget rather than a repartition's + Some(vec![Arc::new(JoinSelection::new())]) + } + } + } +} + +/// Number of distinct string values held by every group produced by +/// [`grouped_distinct_string_batches`] +const DISTINCT_VALUES_PER_GROUP: usize = 2; + +/// Returns batches of 1024 rows with `groups` distinct keys in `group_key`, +/// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings +/// in `value` and an `Int64` `payload` to aggregate over. The values are +/// `Utf8View` if `string_view` is set, `Utf8` otherwise. +fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec { + let value_type = if string_view { + DataType::Utf8View + } else { + DataType::Utf8 + }; + let schema = Arc::new(Schema::new(vec![ + Field::new("group_key", DataType::Int32, false), + Field::new("value", value_type, false), + Field::new("payload", DataType::Int64, false), + ])); + + const ROWS_PER_BATCH: usize = 1024; + + let rows = groups * DISTINCT_VALUES_PER_GROUP; + let mut keys = Vec::with_capacity(rows); + let mut values = Vec::with_capacity(rows); + for value in 0..DISTINCT_VALUES_PER_GROUP { + for group in 0..groups { + keys.push(group as i32); + values.push(format!("value-{value}")); } } + + keys.chunks(ROWS_PER_BATCH) + .zip(values.chunks(ROWS_PER_BATCH)) + .map(|(keys, values)| { + let payloads: ArrayRef = + Arc::new(Int64Array::from_iter_values(keys.iter().map(|k| *k as i64))); + let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec())); + let values: ArrayRef = if string_view { + Arc::new(StringViewArray::from_iter_values(values)) + } else { + Arc::new(StringArray::from_iter_values(values)) + }; + RecordBatch::try_new(Arc::clone(&schema), vec![keys, values, payloads]) + .unwrap() + }) + .collect() } fn access_log_batches() -> Vec { diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 3b708b4edeec9..ad491c9782637 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2514,7 +2514,7 @@ impl NormalizeEq for Expr { args: other_args, }), ) => { - self_func.name() == other_func.name() + self_func == other_func && self_args.len() == other_args.len() && self_args .iter() @@ -2545,7 +2545,7 @@ impl NormalizeEq for Expr { }, }), ) => { - self_func.name() == other_func.name() + self_func == other_func && self_distinct == other_distinct && self_null_treatment == other_null_treatment && self_args.len() == other_args.len() @@ -2598,7 +2598,7 @@ impl NormalizeEq for Expr { }, } = other.as_ref(); - self_fun.name() == other_fun.name() + self_fun == other_fun && self_window_frame == other_window_frame && match (self_filter, other_filter) { (Some(a), Some(b)) => a.normalize_eq(b), @@ -2611,10 +2611,12 @@ impl NormalizeEq for Expr { .iter() .zip(other_args.iter()) .all(|(a, b)| a.normalize_eq(b)) + && self_partition_by.len() == other_partition_by.len() && self_partition_by .iter() .zip(other_partition_by.iter()) .all(|(a, b)| a.normalize_eq(b)) + && self_order_by.len() == other_order_by.len() && self_order_by .iter() .zip(other_order_by.iter()) @@ -3831,6 +3833,7 @@ pub fn physical_name(expr: &Expr) -> Result { #[cfg(test)] mod test { use crate::expr_fn::col; + use crate::test::function_stub::max_udaf; use crate::{ ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Volatility, case, lit, placeholder, qualified_wildcard, wildcard, wildcard_with_options, @@ -4345,6 +4348,32 @@ mod test { use super::*; use crate::logical_plan::{EmptyRelation, LogicalPlan}; + #[test] + fn normalize_eq_window_function_over_clause_lengths() { + let window = |partition_by: Vec, order_by: Vec| { + let mut window = WindowFunction::new(max_udaf(), vec![col("value")]); + window.params.partition_by = partition_by; + window.params.order_by = order_by; + Expr::from(window) + }; + let base = window(vec![col("a")], vec![Sort::new(col("a"), true, true)]); + + let extra_partition = window( + vec![col("a"), col("b")], + vec![Sort::new(col("a"), true, true)], + ); + assert!(!base.normalize_eq(&extra_partition)); + + let extra_order = window( + vec![col("a")], + vec![ + Sort::new(col("a"), true, true), + Sort::new(col("b"), true, true), + ], + ); + assert!(!base.normalize_eq(&extra_order)); + } + #[test] fn test_display_wildcard() { assert_eq!(format!("{}", wildcard()), "*"); diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 36b76f076d26a..75b3a60af1465 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -33,6 +33,7 @@ use crate::{ use arrow::compute::can_cast_types; use arrow::datatypes::FieldRef; use arrow::datatypes::{DataType, Field}; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::FieldExt; use datafusion_common::{ Column, DataFusionError, ExprSchema, Result, ScalarValue, Spans, TableReference, @@ -71,18 +72,46 @@ pub trait ExprSchemable { -> Result<(DataType, bool)>; } -/// Derives the output field for a cast expression from the source field. +/// Derives the output field for a cast expression from the source and target fields. +/// +/// Metadata handling: +/// - Type-only casts (i.e., target_field == DataType::SomeDataType.into_nullable_field()) +/// propagate non extension-type metadata from the source. This is for backward compatibility +/// (casts have propagated source metadata for many if not all previous versions), recognizing +/// that the return type of `::` should have the +/// return type of `` (e.g., casting arrow.json to utf8). +/// - All other casts preserve target metadata exactly. This ensures in particular that output +/// metadata when casting to an extension type contains the extension information in the +/// output field. Callers that wish to have some mix of source and target metadata can use +/// Alias or construct an output field themselves (whose metadata will be used directly). +/// /// For `TryCast`, `force_nullable` is `true` since a failed cast returns NULL. fn cast_output_field( source_field: &FieldRef, - target_type: &DataType, + target_field: &FieldRef, force_nullable: bool, ) -> Arc { + // Check if this is a "type-only" cast (target_field == DataType::X.into_nullable_field()) + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + let metadata = if is_type_only { + // Type-only cast: propagate source metadata, stripping extension type keys + let mut meta = source_field.metadata().clone(); + meta.remove(EXTENSION_TYPE_NAME_KEY); + meta.remove(EXTENSION_TYPE_METADATA_KEY); + meta + } else { + // Explicit target field: use target metadata exactly + target_field.metadata().clone() + }; + let mut f = source_field .as_ref() .clone() - .with_data_type(target_type.clone()) - .with_metadata(source_field.metadata().clone()); + .with_data_type(target_field.data_type().clone()) + .with_metadata(metadata); if force_nullable { f = f.with_nullable(true); } @@ -478,7 +507,8 @@ impl ExprSchemable for Expr { /// - **Aliases**: Merge underlying expr metadata with alias-specific metadata, preferring the alias metadata /// - **Binary expressions**: field metadata is empty /// - **Boolean expressions**: field metadata is empty - /// - **Cast expressions**: determined by the input expression's field metadata handling + /// - **Cast expressions**: Type-only casts pass through source metadata (stripping extension + /// type keys); casts with explicit target fields use target metadata exactly /// - **Scalar functions**: Generate metadata via function's [`return_field_from_args`] method, /// with the default implementation returning empty field metadata /// - **Aggregate functions**: Generate metadata via function's [`return_field`] method, @@ -623,20 +653,16 @@ impl ExprSchemable for Expr { func.return_field_from_args(args) } // _ => Ok((self.get_type(schema)?, self.nullable(schema)?)), - Expr::Cast(Cast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), false) - }) - } + Expr::Cast(Cast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, false)), Expr::Placeholder(Placeholder { id: _, field: Some(field), }) => Ok(Arc::clone(field).renamed(&schema_name)), - Expr::TryCast(TryCast { expr, field }) => { - expr.to_field(schema).map(|(_table_ref, src)| { - cast_output_field(&src, field.data_type(), true) - }) - } + Expr::TryCast(TryCast { expr, field }) => expr + .to_field(schema) + .map(|(_table_ref, src)| cast_output_field(&src, field, true)), Expr::LambdaVariable(LambdaVariable { field: Some(field), .. }) => Ok(Arc::clone(field).renamed(&schema_name)), @@ -1427,4 +1453,115 @@ mod tests { assert_eq!(meta, expr.metadata(&schema).unwrap()); } + + #[test] + fn test_cast_and_try_cast_extension_type_metadata() { + use crate::expr::{Cast, TryCast}; + use arrow_schema::extension::{ + EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY, + }; + + // Helper to build either Cast or TryCast expression + fn make_cast_expr( + expr: Expr, + target_field: FieldRef, + use_try_cast: bool, + ) -> Expr { + if use_try_cast { + Expr::TryCast(TryCast { + expr: Box::new(expr), + field: target_field, + }) + } else { + Expr::Cast(Cast { + expr: Box::new(expr), + field: target_field, + }) + } + } + + // Run the same test logic for both Cast and TryCast + for use_try_cast in [false, true] { + let cast_name = if use_try_cast { "TryCast" } else { "Cast" }; + + // Create a schema with a field that has extension type metadata + let mut source_meta = HashMap::new(); + source_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ); + source_meta.insert("custom_key".to_string(), "custom_value".to_string()); + + let source_field = Field::new("foo", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta); + + let schema = MockExprSchema::new() + .with_data_type(DataType::FixedSizeBinary(16)) + .with_metadata(FieldMetadata::from(source_field.metadata().clone())); + + // Test 1: Cast to a type without extension metadata strips extension metadata + // but preserves non-extension metadata + let cast_expr = make_cast_expr( + col("foo"), + Arc::new(Field::new("", DataType::Utf8, true)), + use_try_cast, + ); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert!( + result_field + .metadata() + .get(EXTENSION_TYPE_NAME_KEY) + .is_none(), + "{cast_name}: Extension type name should be stripped when target has no extension metadata" + ); + assert_eq!( + result_field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "{cast_name}: Non-extension metadata should be preserved" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + + // Test 2: Cast to a field with explicit metadata uses target metadata exactly + let mut target_meta = HashMap::new(); + target_meta.insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ); + target_meta.insert(EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()); + + let target_field = + Field::new("", DataType::Utf8, true).with_metadata(target_meta); + + let cast_expr = + make_cast_expr(col("foo"), Arc::new(target_field), use_try_cast); + + let (_, result_field) = cast_expr.to_field(&schema).unwrap(); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()), + "{cast_name}: Extension type name should come from target field" + ); + assert_eq!( + result_field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()), + "{cast_name}: Extension type metadata should come from target field" + ); + assert!( + result_field.metadata().get("custom_key").is_none(), + "{cast_name}: Source metadata should NOT propagate when target has explicit metadata" + ); + if use_try_cast { + assert!( + result_field.is_nullable(), + "TryCast result should be nullable" + ); + } + } + } } diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 781559ddd0c5c..8e86cb3685e90 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -1158,6 +1158,16 @@ fn coerced_from<'a>( { Some(type_into.clone()) } + (_, RunEndEncoded(_, value_type)) + if coerced_from(type_into, value_type.data_type()).is_some() => + { + Some(type_into.clone()) + } + (RunEndEncoded(_, value_type), _) + if coerced_from(value_type.data_type(), type_from).is_some() => + { + Some(type_into.clone()) + } // coerced into type_into (Int8, Null | Int8) => Some(type_into.clone()), (Int16, Null | Int8 | Int16 | UInt8) => Some(type_into.clone()), @@ -1617,6 +1627,50 @@ mod tests { ); } + #[test] + fn test_coerced_from_run_end_encoded() { + let run_end_encoded_of = |value_type: DataType| { + DataType::RunEndEncoded( + Field::new("run_ends", DataType::Int32, false).into(), + Field::new("values", value_type, true).into(), + ) + }; + + let type_into = run_end_encoded_of(DataType::UInt32); + let type_from = DataType::Int64; + assert_eq!(coerced_from(&type_into, &type_from), None); + + let type_from = run_end_encoded_of(DataType::UInt32); + let type_into = DataType::Int64; + assert_eq!( + coerced_from(&type_into, &type_from), + Some(type_into.clone()) + ); + + // Signature candidates for functions like `date_bin` are plain + // Timestamp, but a REE-encoded column (e.g. a segment written with + // REE-dict encoding for that field) should still coerce against + // them via the wrapped value type. + let type_from = + run_end_encoded_of(DataType::Timestamp(TimeUnit::Nanosecond, None)); + let type_into = DataType::Timestamp(TimeUnit::Nanosecond, None); + assert_eq!( + coerced_from(&type_into, &type_from), + Some(type_into.clone()) + ); + + // The reverse direction: a plain type coercing into an REE target + // (e.g. a signature that happens to require RunEndEncoded) should + // succeed whenever the plain type coerces into the wrapped value + // type. + let type_into = run_end_encoded_of(DataType::Int64); + let type_from = DataType::Int32; + assert_eq!( + coerced_from(&type_into, &type_from), + Some(type_into.clone()) + ); + } + #[test] fn test_get_valid_types_array_and_array() -> Result<()> { let function = "array_and_array"; diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs index f6df4182a879b..d955d343ad629 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs @@ -39,6 +39,18 @@ use std::mem::size_of_val; pub struct BytesDistinctCountAccumulator(ArrowBytesSet); impl BytesDistinctCountAccumulator { + /// The set deliberately does not pre-allocate. `GroupsAccumulatorAdapter` + /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a + /// high cardinality key holds hundreds of thousands of these at once and + /// most of them see only a handful of values. + /// + /// The ungrouped path builds one of these and grows it to hold every + /// distinct value in the input, so it is the caller that had a use for the + /// warm up. It loses nothing here: the set grows into exactly the + /// capacities a pre-allocated one reaches, which + /// `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set` + /// pins. That is why this constructor needs no signal distinguishing the + /// two callers. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesSet::new(output_type)) } @@ -100,6 +112,8 @@ impl Accumulator for BytesDistinctCountAccumulator { pub struct BytesViewDistinctCountAccumulator(ArrowBytesViewSet); impl BytesViewDistinctCountAccumulator { + /// See [`BytesDistinctCountAccumulator::new`] for why the set does not + /// pre-allocate. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesViewSet::new(output_type)) } @@ -151,3 +165,144 @@ impl Accumulator for BytesViewDistinctCountAccumulator { size_of_val(self) + self.0.size() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{StringArray, StringViewArray}; + use datafusion_physical_expr_common::binary_map::INITIAL_MAP_CAPACITY; + use datafusion_physical_expr_common::binary_view_map::INITIAL_MAP_CAPACITY as INITIAL_VIEW_MAP_CAPACITY; + use std::sync::Arc; + + /// The batch size the aggregate stream drives an ungrouped accumulator with. + const BATCH_SIZE: usize = 8192; + + /// Distinct value counts spanning both sides of the warm up capacity, up to + /// ones where the two constructors have converged. + const CARDINALITIES: [usize; 7] = [0, 1, 100, 1_000, 10_000, 100_000, 500_000]; + + /// Cardinalities small enough that the warm up dominates what the set + /// holds. This is the per group population, where `GroupsAccumulatorAdapter` + /// holds one accumulator per group and most see a handful of values. + const PER_GROUP_SCALE: usize = 100; + + /// Cardinalities at which a lazily built set has grown into exactly the + /// capacities a pre-allocated one reaches. This is the ungrouped + /// population, one set holding every distinct value in the input. + const UNGROUPED_SCALE: usize = 10_000; + + /// Longer than the map's inline value length, so the value lands in the + /// value buffer rather than inside the hash table entry. + fn distinct_value(i: usize) -> String { + format!("distinct value number {i}") + } + + fn batches(distinct_values: usize, view: bool) -> Vec { + (0..distinct_values) + .step_by(BATCH_SIZE) + .map(|start| { + let values = (start..(start + BATCH_SIZE).min(distinct_values)) + .map(distinct_value); + if view { + Arc::new(StringViewArray::from_iter_values(values)) as ArrayRef + } else { + Arc::new(StringArray::from_iter_values(values)) as ArrayRef + } + }) + .collect() + } + + /// The property that decides whether the ungrouped path can afford to share + /// the lazy constructor with the per group path. + fn assert_lazy_is_not_worse( + distinct_values: usize, + lazy_size: usize, + pre_allocated_size: usize, + ) { + // The guarantee that lets both paths share one lazy constructor. + assert!( + lazy_size <= pre_allocated_size, + "at {distinct_values} distinct values the lazy set reported \ + {lazy_size} bytes against the {pre_allocated_size} bytes the \ + pre-allocated one reported" + ); + + if distinct_values <= PER_GROUP_SCALE { + // The warm up is pure overhead here, and removing it is the whole + // point of the change. + assert!( + lazy_size < pre_allocated_size, + "at {distinct_values} distinct values the lazy set should be \ + strictly cheaper, but reported {lazy_size} bytes against \ + {pre_allocated_size}" + ); + } + + if distinct_values >= UNGROUPED_SCALE { + // The hash table's bucket count is a power of two fixed by the + // number of entries, and the value buffer grows on a power of two + // ladder, so a set that starts empty lands on exactly the + // capacities a pre-allocated one reaches. The ungrouped path gives + // up nothing by starting empty, which is why these accumulators + // need no signal telling them apart from the per group ones. + assert_eq!( + lazy_size, pre_allocated_size, + "at {distinct_values} distinct values the lazy and pre-allocated \ + sets should have converged" + ); + } + } + + /// An ungrouped `COUNT(DISTINCT )` builds a single accumulator that + /// grows to hold every distinct value in its input, which is the population + /// the warm up existed for. It must be no more expensive without one. + #[test] + fn ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set() { + for distinct_values in CARDINALITIES { + let mut accumulator = + BytesDistinctCountAccumulator::::new(OutputType::Utf8); + let mut pre_allocated = ArrowBytesSet::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + for batch in batches(distinct_values, false) { + accumulator.update_batch(&[Arc::clone(&batch)]).unwrap(); + pre_allocated.insert(&batch); + } + + assert_eq!(accumulator.0.non_null_len(), distinct_values); + assert_lazy_is_not_worse( + distinct_values, + accumulator.0.size(), + pre_allocated.size(), + ); + } + } + + /// The `Utf8View` counterpart of + /// [`ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`]. + #[test] + fn ungrouped_utf8_view_accumulator_is_never_worse_than_a_pre_allocated_set() { + for distinct_values in CARDINALITIES { + let mut accumulator = + BytesViewDistinctCountAccumulator::new(OutputType::Utf8View); + let mut pre_allocated = ArrowBytesViewSet::with_capacity( + OutputType::Utf8View, + INITIAL_VIEW_MAP_CAPACITY, + ); + + for batch in batches(distinct_values, true) { + accumulator.update_batch(&[Arc::clone(&batch)]).unwrap(); + pre_allocated.insert(&batch); + } + + assert_eq!(accumulator.0.non_null_len(), distinct_values); + assert_lazy_is_not_worse( + distinct_values, + accumulator.0.size(), + pre_allocated.size(), + ); + } + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index 6e3e3b91a74f7..38f6d26200ec9 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -89,7 +89,10 @@ where match emit_to { EmitTo::All => { - self.seen.clear(); + // Release the capacity, not just the entries: `size()` reports + // capacity, and the aggregate streams rely on it dropping after + // emitting everything. + self.seen = HashSet::default(); } EmitTo::First(n) => { let mut remaining = HashSet::default(); @@ -145,7 +148,9 @@ where all_values[pos] = value; cursors[group_idx] += 1; } - self.counts.clear(); + // Release the capacity, see `evaluate`. + self.seen = HashSet::default(); + self.counts = Vec::new(); } else { let mut remaining = HashSet::default(); for (group_idx, value) in self.seen.drain() { diff --git a/datafusion/functions-nested/src/position.rs b/datafusion/functions-nested/src/position.rs index f677cfb4979d3..9ef8142fa48dd 100644 --- a/datafusion/functions-nested/src/position.rs +++ b/datafusion/functions-nested/src/position.rs @@ -40,7 +40,9 @@ use arrow::array::{ use datafusion_common::cast::{ as_generic_list_array, as_int64_array, as_large_list_array, as_list_array, }; -use datafusion_common::{Result, exec_err, utils::take_function_args}; +use datafusion_common::{ + Result, exec_datafusion_err, exec_err, utils::take_function_args, +}; use itertools::Itertools; use crate::utils::{compare_element_to_list, make_scalar_function}; @@ -198,6 +200,16 @@ fn array_position_inner(args: &[ArrayRef]) -> Result { } } +fn resolve_zero_based_start_from(start_from: i64) -> Result { + start_from.checked_sub(1).ok_or_else(|| { + exec_datafusion_err!( + "start_from out of bounds: {start_from}, expected {} to {}", + i64::MIN + 1, + i64::MAX + ) + }) +} + /// Resolves the optional `start_from` argument into a `Vec` of /// 0-indexed starting positions. fn resolve_start_from( @@ -207,14 +219,16 @@ fn resolve_start_from( match third_arg { None => Ok(vec![0i64; num_rows]), Some(ColumnarValue::Scalar(ScalarValue::Int64(Some(v)))) => { - Ok(vec![v - 1; num_rows]) + Ok(vec![resolve_zero_based_start_from(*v)?; num_rows]) } Some(ColumnarValue::Scalar(s)) => { exec_err!("array_position expected Int64 for start_from, got {s}") } - Some(ColumnarValue::Array(a)) => { - Ok(as_int64_array(a)?.values().iter().map(|&x| x - 1).collect()) - } + Some(ColumnarValue::Array(a)) => as_int64_array(a)? + .values() + .iter() + .map(|&x| resolve_zero_based_start_from(x)) + .collect(), } } @@ -309,8 +323,8 @@ fn general_position_dispatch(args: &[ArrayRef]) -> Result>() + .map(|&x| resolve_zero_based_start_from(x)) + .collect::>>()? } else { vec![0; haystack.len()] }; diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index 00d1531a3a771..27c644042becb 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -68,7 +68,7 @@ name = "datafusion_functions" arrow = { workspace = true } arrow-buffer = { workspace = true } base64 = { version = "0.23", optional = true } -blake2 = { version = "^0.10.2", optional = true } +blake2 = { version = "^0.11.0", optional = true } blake3 = { version = "1.8", optional = true } chrono = { workspace = true } chrono-tz = { version = "0.10.4", optional = true } diff --git a/datafusion/functions/src/core/arrow_cast.rs b/datafusion/functions/src/core/arrow_cast.rs index 0b67883c17c87..929a92abcca23 100644 --- a/datafusion/functions/src/core/arrow_cast.rs +++ b/datafusion/functions/src/core/arrow_cast.rs @@ -27,8 +27,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -161,9 +161,27 @@ impl ScalarUDFImpl for ArrowCastFunc { let [source_arg, type_arg] = take_function_args(self.name(), args)?; let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; + + // We can skip the cast only if: + // 1. The source and target types are the same + // 2. The source has no extension metadata that needs to be stripped let new_expr = if source_type == target_type { - // the argument's data type is already the correct type - source_arg + // Check if source has extension metadata + let source_field = source_arg.to_field(info.schema())?; + let has_extension_metadata = source_field + .1 + .metadata() + .contains_key("ARROW:extension:name"); + if has_extension_metadata { + // Need to create a cast to strip extension metadata + Expr::Cast(datafusion_expr::Cast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + // the argument's data type is already the correct type + source_arg + } } else { // Use an actual cast to get the correct type Expr::Cast(datafusion_expr::Cast { diff --git a/datafusion/functions/src/core/arrow_try_cast.rs b/datafusion/functions/src/core/arrow_try_cast.rs index d27b29ba5736d..0914695b60370 100644 --- a/datafusion/functions/src/core/arrow_try_cast.rs +++ b/datafusion/functions/src/core/arrow_try_cast.rs @@ -26,8 +26,8 @@ use datafusion_common::{ use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ReturnFieldArgs, ScalarFunctionArgs, - ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, Expr, ExprSchemable, ReturnFieldArgs, + ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -134,8 +134,26 @@ impl ScalarUDFImpl for ArrowTryCastFunc { let target_type = data_type_from_type_arg(self.name(), &type_arg)?; let source_type = info.get_data_type(&source_arg)?; + + // We can skip the cast only if: + // 1. The source and target types are the same + // 2. The source has no extension metadata that needs to be stripped let new_expr = if source_type == target_type { - source_arg + // Check if source has extension metadata + let source_field = source_arg.to_field(info.schema())?; + let has_extension_metadata = source_field + .1 + .metadata() + .contains_key("ARROW:extension:name"); + if has_extension_metadata { + // Need to create a try_cast to strip extension metadata + Expr::TryCast(datafusion_expr::TryCast { + expr: Box::new(source_arg), + field: target_type.into_nullable_field_ref(), + }) + } else { + source_arg + } } else { Expr::TryCast(datafusion_expr::TryCast { expr: Box::new(source_arg), diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 0f82f0b0df764..13b9541b6a71a 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -1088,6 +1088,8 @@ fn extract_window_frame_target_type(col_type: &DataType) -> Result { Ok(DataType::Interval(IntervalUnit::MonthDayNano)) } else if let DataType::Dictionary(_, value_type) = col_type { extract_window_frame_target_type(value_type) + } else if let DataType::RunEndEncoded(_, value_type) = col_type { + extract_window_frame_target_type(value_type.data_type()) } else { internal_err!("Cannot run range queries on datatype: {col_type}") } @@ -1113,8 +1115,11 @@ fn coerce_window_frame( // `current_value ± offset`, so it is only meaningful for target // types that support arithmetic. Other orderable target types can // still use free range frames, whose bounds require comparison only. + // REE arrays are not supported by arrow's numeric kernesl. + // Tracked at https://github.com/apache/arrow-rs/issues/10891). let supports_offset_arithmetic = - target_type.is_numeric() || is_interval(&target_type); + !matches!(col_type, DataType::RunEndEncoded(_, _)) + && (target_type.is_numeric() || is_interval(&target_type)); if !supports_offset_arithmetic && !window_frame.free_range() { return plan_err!( "RANGE with offset PRECEDING/FOLLOWING is not supported for ORDER BY type {target_type}" diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index 107758cf6b6e7..dfaf93bc1e8fa 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -57,6 +57,58 @@ fn has_extractable_expr(exprs: &[Expr]) -> bool { }) } +/// Returns the flat names of `plan`'s output columns whose defining expression +/// is volatile (e.g. `random()`). +/// +/// Only a [`LogicalPlan::Projection`] can define such a column: anywhere else +/// the value has already been materialized by the projection that produced it, +/// so referencing the column again does not re-evaluate anything. +fn volatile_output_columns(plan: &LogicalPlan) -> BTreeSet { + let LogicalPlan::Projection(projection) = plan else { + return BTreeSet::new(); + }; + projection + .schema + .iter() + .zip(projection.expr.iter()) + .filter(|(_, expr)| expr.is_volatile()) + .map(|((qualifier, field), _)| Column::from((qualifier, field)).flat_name()) + .collect() +} + +/// Returns `true` if building an extraction projection for `exprs` on top of +/// `input` would duplicate a volatile computation. +/// +/// When `input` is already a projection, [`build_extraction_projection_impl`] +/// *merges* into it: every column reference in an extracted expression is +/// replaced by that column's defining expression (see +/// [`build_projection_replace_map`]). Inlining a volatile definition makes the +/// merged projection evaluate it a second, independent time, so the extracted +/// value no longer matches the column it was derived from: +/// +/// ```text +/// Projection: s, get_field(s, 'a') AS field +/// Projection: named_struct('a', random()) AS s +/// ``` +/// +/// would merge into a single projection computing `random()` twice, and +/// `field` would then differ from `s['a']` on every row. Callers skip the +/// extraction instead. +fn would_duplicate_volatile<'a>( + exprs: impl IntoIterator, + input: &LogicalPlan, +) -> bool { + let volatile = volatile_output_columns(input); + if volatile.is_empty() { + return false; + } + exprs.into_iter().any(|expr| { + expr.column_refs() + .iter() + .any(|col| volatile.contains(&col.flat_name())) + }) +} + /// Extracts `MoveTowardsLeafNodes` sub-expressions from non-projection nodes /// into **extraction projections** (pass 1 of 2). /// @@ -195,7 +247,18 @@ fn extract_from_plan( } // Fast pre-check: skip all allocations if no extractable expressions exist - if !has_extractable_expr(&plan.expressions()) { + let node_exprs = plan.expressions(); + if !has_extractable_expr(&node_exprs) { + return Ok(Transformed::no(plan)); + } + + // The extraction projection is merged into an input that is already a + // projection, which inlines the referenced columns' definitions. Skip the + // extraction when that would duplicate a volatile computation. + if inputs + .iter() + .any(|input| would_duplicate_volatile(node_exprs.iter(), input)) + { return Ok(Transformed::no(plan)); } @@ -917,6 +980,16 @@ fn split_and_push_projection( return Ok(None); } + // Pushing into an input that is already a projection merges into it and + // inlines the referenced columns' definitions. Leave the projection alone + // when that would duplicate a volatile computation. + if would_duplicate_volatile( + extraction_pairs.iter().map(|(expr, _)| expr), + input.as_ref(), + ) { + return Ok(None); + } + // ── Phase 2: Push down ────────────────────────────────────────────── let proj_input = Arc::clone(&proj.input); let pushed = push_extraction_pairs( @@ -1215,6 +1288,15 @@ fn try_push_into_inputs( if per_input[idx].pairs.is_empty() { new_inputs.push(input.clone()); } else { + // Merging into an input projection inlines the referenced columns' + // definitions; bail out when that would duplicate a volatile + // computation. + if would_duplicate_volatile( + per_input[idx].pairs.iter().map(|(expr, _)| expr), + input, + ) { + return Ok(None); + } let input_arc = Arc::new(input.clone()); let target_schema = Arc::clone(input.schema()); let proj = build_extraction_projection_impl( @@ -1270,7 +1352,7 @@ mod tests { use crate::{Optimizer, OptimizerContext}; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{ - ScalarUDF, col, lit, logical_plan::builder::LogicalPlanBuilder, + ScalarUDF, Volatility, col, lit, logical_plan::builder::LogicalPlanBuilder, }; fn leaf_udf(expr: Expr, name: &str) -> Expr { @@ -1283,6 +1365,19 @@ mod tests { )) } + /// A stand-in for `random()`: a volatile expression that must be evaluated + /// exactly once per row. + fn volatile_udf(expr: Expr) -> Expr { + Expr::ScalarFunction(ScalarFunction::new_udf( + Arc::new(ScalarUDF::new_from_impl( + PlacementTestUDF::new() + .with_placement(ExpressionPlacement::KeepInPlace) + .with_volatility(Volatility::Volatile), + )), + vec![expr], + )) + } + // ========================================================================= // Combined optimization stage formatter // ========================================================================= @@ -1960,6 +2055,63 @@ mod tests { "#) } + /// Merging an extraction into an input projection inlines the definition of + /// every column the extraction references. When that definition is volatile + /// the inlined copy is an independent evaluation, so the extraction must be + /// skipped and the plan left alone. + #[test] + fn test_no_merge_into_volatile_projection() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .project(vec![volatile_udf(col("a")).alias("v")])? + .project(vec![col("v"), leaf_udf(col("v"), "x")])? + .build()?; + + assert_stages!(plan, @r#" + ## Original Plan + Projection: v, leaf_udf(v, Utf8("x")) + Projection: keep_in_place_udf(test.a) AS v + TableScan: test projection=[a] + + ## After Extraction + (same as original) + + ## After Pushdown + (same as after extraction) + + ## Optimized + (same as after pushdown) + "#) + } + + /// Same guard for pass 1: extracting out of a `Filter` whose input + /// projection defines the referenced column volatilely would make the + /// predicate test a second, independent evaluation. + #[test] + fn test_no_extraction_from_filter_over_volatile_projection() -> Result<()> { + let table_scan = test_table_scan()?; + let plan = LogicalPlanBuilder::from(table_scan) + .project(vec![volatile_udf(col("a")).alias("v")])? + .filter(leaf_udf(col("v"), "x").eq(lit(1u32)))? + .build()?; + + assert_stages!(plan, @r#" + ## Original Plan + Filter: leaf_udf(v, Utf8("x")) = UInt32(1) + Projection: keep_in_place_udf(test.a) AS v + TableScan: test projection=[a] + + ## After Extraction + (same as original) + + ## After Pushdown + (same as after extraction) + + ## Optimized + (same as after pushdown) + "#) + } + /// Projections with aliased columns (nothing to extract) return unchanged. #[test] fn test_projection_early_return_no_extraction() -> Result<()> { diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index b0543a871f52b..9604bf2f188c0 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -556,17 +556,6 @@ fn merge_consecutive_projections_one_level( return Projection::try_new_with_schema(expr, input, schema).map(Transformed::no); }; - // A fast path: if the previous projection is same as the current projection - // we can directly remove the current projection and return child projection. - if prev_projection.expr == expr { - return Projection::try_new_with_schema( - expr, - Arc::clone(&prev_projection.input), - schema, - ) - .map(Transformed::yes); - } - // Count usages (referrals) of each projection expression in its input fields: let mut column_referral_map = HashMap::<&Column, usize>::new(); for expr in &expr { @@ -1210,6 +1199,69 @@ mod tests { ) } + #[test] + fn merge_structurally_equal_non_idempotent_projections() -> Result<()> { + let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]); + let projection = || col("i").add(lit(1)).alias("i"); + let plan = table_scan(TableReference::none(), &schema, None)? + .project(vec![projection()])? + .project(vec![projection()])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: ?table?.i + Int32(1) + Int32(1) AS i + TableScan: ?table? projection=[i] + " + ) + } + + #[test] + fn merge_deep_projection_chain_in_one_pass() -> Result<()> { + let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]); + let mut plan = table_scan(TableReference::none(), &schema, None)? + .project(vec![col("i").add(lit(1)).alias("i")])? + .build()?; + for _ in 1..12 { + plan = LogicalPlanBuilder::from(plan) + .project(vec![col("i").add(lit(1)).alias("i")])? + .build()?; + } + + let optimizer = Optimizer::with_rules(vec![Arc::new(OptimizeProjections::new())]); + let optimized = optimizer.optimize( + plan, + &OptimizerContext::new().with_max_passes(1), + observe, + )?; + let plan_string = format!("{optimized}"); + assert_eq!(12, plan_string.matches("Int32(1)").count()); + assert_eq!(1, plan_string.matches("Projection:").count()); + Ok(()) + } + + #[test] + fn merge_columns_and_metadata_alias() -> Result<()> { + let metadata = + datafusion_common::metadata::FieldMetadata::from(HashMap::from([( + "key".to_string(), + "value".to_string(), + )])); + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .project(vec![col("a")])? + .project(vec![col("a").alias_with_metadata("a", Some(metadata))])? + .build()?; + + let optimized = optimize(plan)?; + assert_eq!( + "value", + optimized.schema().field(0).metadata().get("key").unwrap() + ); + assert_eq!(1, format!("{optimized}").matches("Projection:").count()); + Ok(()) + } + #[test] fn merge_alias() -> Result<()> { let table_scan = test_table_scan()?; diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index dadea4784802a..ec01becb8f7a5 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -2514,6 +2514,77 @@ mod tests { } } + #[test] + fn test_simplify_swapped_operands_in_and_or_no_canonicalize() { + // Regression test for https://github.com/apache/datafusion/issues/14943 + // + // `SimplifyExpressions` disables canonicalization for `LogicalPlan::Join` + // (see https://github.com/apache/datafusion/pull/8780), so commutative + // operands like `A = B` and `B = A` cannot be normalized to a single + // form before AND/OR dedup runs. The dedup itself must therefore + // recognize commutative equivalence directly. + + // c3 = 5 AND 5 = c3 --> c3 = 5 + let expr = col("c3_non_null") + .eq(lit(5_i64)) + .and(lit(5_i64).eq(col("c3_non_null"))); + let expected = col("c3_non_null").eq(lit(5_i64)); + assert_eq!(simplify_no_canonicalize(expr), expected); + + // 5 = c3 AND c3 = 5 --> 5 = c3 + let expr = lit(5_i64) + .eq(col("c3_non_null")) + .and(col("c3_non_null").eq(lit(5_i64))); + let expected = lit(5_i64).eq(col("c3_non_null")); + assert_eq!(simplify_no_canonicalize(expr), expected); + + // c3 = 5 OR 5 = c3 --> c3 = 5 + let expr = col("c3_non_null") + .eq(lit(5_i64)) + .or(lit(5_i64).eq(col("c3_non_null"))); + let expected = col("c3_non_null").eq(lit(5_i64)); + assert_eq!(simplify_no_canonicalize(expr), expected); + + // (c3 = 5 AND c4 > 0) AND (5 = c3) --> c3 = 5 AND c4 > 0 + let expr = col("c3_non_null") + .eq(lit(5_i64)) + .and(col("c4_non_null").gt(lit(0_u32))) + .and(lit(5_i64).eq(col("c3_non_null"))); + let expected = col("c3_non_null") + .eq(lit(5_i64)) + .and(col("c4_non_null").gt(lit(0_u32))); + assert_eq!(simplify_no_canonicalize(expr), expected); + } + + #[test] + fn test_simplify_swapped_operands_in_xor_no_canonicalize() { + // `expr_contains` guards the XOR rules, so `delete_xor_in_complex_expr` has to + // use the same equality relation. Comparing structurally there while the guard + // normalizes makes the rule rebuild its input and still report `Transformed::yes`, + // which spins the simplifier until it hits the cycle limit. + // + // Operands are non-nullable so the expected results do not depend on how the XOR + // rules treat a NULL operand, which is a separate question from the cycle. + + // (c4_non_null + 1) ^ (1 + c4_non_null) --> 0 + let expr = bitwise_xor( + col("c4_non_null") + lit(1_u32), + lit(1_u32) + col("c4_non_null"), + ); + let (simplified, cycles) = simplify_no_canonicalize_with_cycle_count(expr); + assert_eq!(simplified, lit(0_u32)); + assert_eq!(cycles, 2); + + // (c4_non_null + 1) ^ ((1 + c4_non_null) ^ c4_non_null) --> c4_non_null + let expr = bitwise_xor( + col("c4_non_null") + lit(1_u32), + bitwise_xor(lit(1_u32) + col("c4_non_null"), col("c4_non_null")), + ); + let (simplified, cycles) = simplify_no_canonicalize_with_cycle_count(expr); + assert_eq!(simplified, col("c4_non_null")); + assert_eq!(cycles, 2); + } + #[test] fn test_simplify_eq_not_self() { // `expr_a`: column `c2` is nullable, so `c2 = c2` simplifies to `c2 IS NOT NULL OR NULL` @@ -3667,6 +3738,20 @@ mod tests { try_simplify(expr).unwrap() } + fn simplify_no_canonicalize(expr: Expr) -> Expr { + simplify_no_canonicalize_with_cycle_count(expr).0 + } + + fn simplify_no_canonicalize_with_cycle_count(expr: Expr) -> (Expr, u32) { + let schema = expr_test_schema(); + let (transformed, count) = + ExprSimplifier::new(SimplifyContext::builder().with_schema(schema).build()) + .with_canonicalize(false) + .simplify_with_cycle_count_transformed(expr) + .unwrap(); + (transformed.data, count) + } + fn try_simplify_with_cycle_count(expr: Expr) -> Result<(Expr, u32)> { let schema = expr_test_schema(); let simplifier = diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs b/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs index 356f2711b708e..e7edc34cfe4e6 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs @@ -52,27 +52,45 @@ pub fn simplify_predicates(predicates: Vec) -> Result> { let mut other_predicates = Vec::new(); for pred in predicates { - match &pred { - Expr::BinaryExpr(BinaryExpr { - left, - op: + match pred { + Expr::BinaryExpr(BinaryExpr { left, op, right }) + if matches!( + op, Operator::Gt - | Operator::GtEq - | Operator::Lt - | Operator::LtEq - | Operator::Eq, - right, - }) => { + | Operator::GtEq + | Operator::Lt + | Operator::LtEq + | Operator::Eq + ) => + { if let (Some(col), Some(_)) = - (extract_column_from_expr(left), right.as_literal()) - { - column_predicates.entry(col).or_default().push(pred); - } else if let (Some(_), Some(col)) = - (left.as_literal(), extract_column_from_expr(right)) + (extract_column_from_expr(&left), right.as_literal()) { - column_predicates.entry(col).or_default().push(pred); + column_predicates + .entry(col) + .or_default() + .push(Expr::BinaryExpr(BinaryExpr { left, op, right })); + } else if let (Some(_), Some(col), Some(swapped_op)) = ( + left.as_literal(), + extract_column_from_expr(&right), + op.swap(), + ) { + // Put the literal on the right so that predicates differing only in + // operand order compare equal below + column_predicates + .entry(col) + .or_default() + .push(Expr::BinaryExpr(BinaryExpr { + left: right, + op: swapped_op, + right: left, + })); } else { - other_predicates.push(pred); + other_predicates.push(Expr::BinaryExpr(BinaryExpr { + left, + op, + right, + })); } } _ => other_predicates.push(pred), @@ -114,20 +132,12 @@ fn simplify_column_predicates(predicates: Vec) -> Result> { for pred in predicates { match &pred { - Expr::BinaryExpr(BinaryExpr { left: _, op, right }) => { - match (op, right.as_literal().is_some()) { - (Operator::Gt, true) - | (Operator::Lt, false) - | (Operator::GtEq, true) - | (Operator::LtEq, false) => greater_predicates.push(pred), - (Operator::Lt, true) - | (Operator::Gt, false) - | (Operator::LtEq, true) - | (Operator::GtEq, false) => less_predicates.push(pred), - (Operator::Eq, _) => eq_predicates.push(pred), - _ => unreachable!("Unexpected operator: {}", op), - } - } + Expr::BinaryExpr(BinaryExpr { op, .. }) => match op { + Operator::Gt | Operator::GtEq => greater_predicates.push(pred), + Operator::Lt | Operator::LtEq => less_predicates.push(pred), + Operator::Eq => eq_predicates.push(pred), + _ => unreachable!("Unexpected operator: {}", op), + }, _ => unreachable!("Unexpected predicate {}", pred.to_string()), } } @@ -305,6 +315,26 @@ mod tests { assert_eq!(extract_column_from_expr(&col_expr), Some(Column::from("a"))); } + #[test] + fn test_eq_predicates_are_matched_regardless_of_operand_order() { + // a = 5 AND 5 = a is a single condition, not a contradiction + let predicates = vec![col("a").eq(lit(5i32)), lit(5i32).eq(col("a"))]; + + let result = simplify_predicates(predicates).unwrap(); + + assert_eq!(result, vec![col("a").eq(lit(5i32))]); + } + + #[test] + fn test_strict_bound_wins_when_literal_is_on_the_left() { + // a >= 5 AND 5 < a is a > 5; keeping a >= 5 would let a = 5 through + let predicates = vec![col("a").gt_eq(lit(5i32)), lit(5i32).lt(col("a"))]; + + let result = simplify_predicates(predicates).unwrap(); + + assert_eq!(result, vec![col("a").gt(lit(5i32))]); + } + #[test] fn test_simplify_predicates_direct_columns_only() { // Test that only predicates on direct columns are simplified together diff --git a/datafusion/optimizer/src/simplify_expressions/utils.rs b/datafusion/optimizer/src/simplify_expressions/utils.rs index 89bb762d59ce2..78d801630c7ba 100644 --- a/datafusion/optimizer/src/simplify_expressions/utils.rs +++ b/datafusion/optimizer/src/simplify_expressions/utils.rs @@ -17,6 +17,7 @@ //! Utility functions for expression simplification +use datafusion_common::cse::NormalizeEq; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::{ Case, Expr, Like, Operator, @@ -26,13 +27,18 @@ use datafusion_expr::{ /// returns true if `needle` is found in a chain of search_op /// expressions. Such as: (A AND B) AND C +/// +/// Equality uses [`NormalizeEq`] so commutative operands (`A = B` vs `B = A`) +/// are recognized as the same predicate even when the canonicalizer is +/// disabled, notably for `LogicalPlan::Join` filters (see +/// ). fn expr_contains_inner(expr: &Expr, needle: &Expr, search_op: Operator) -> bool { match expr { Expr::BinaryExpr(BinaryExpr { left, op, right }) if *op == search_op => { expr_contains_inner(left, needle, search_op) || expr_contains_inner(right, needle, search_op) } - _ => expr == needle, + _ => expr.normalize_eq(needle), } } @@ -43,6 +49,12 @@ pub fn expr_contains(expr: &Expr, needle: &Expr, search_op: Operator) -> bool { /// Deletes all 'needles' or remains one 'needle' that are found in a chain of xor /// expressions. Such as: A ^ (A ^ (B ^ A)) +/// +/// Matching uses [`NormalizeEq`] to stay consistent with the [`expr_contains`] +/// guard on the XOR rules. A structural comparison here would let the guard fire +/// on operands this function then fails to delete, so the rule would rebuild its +/// input and still report a transformation, spinning the simplifier until it hits +/// the cycle limit. pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> Expr { /// Deletes recursively 'needles' in a chain of xor expressions fn recursive_delete_xor_in_expr( @@ -56,10 +68,10 @@ pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> { let left_expr = recursive_delete_xor_in_expr(left, needle, xor_counter); let right_expr = recursive_delete_xor_in_expr(right, needle, xor_counter); - if left_expr == *needle { + if left_expr.normalize_eq(needle) { *xor_counter += 1; return right_expr; - } else if right_expr == *needle { + } else if right_expr.normalize_eq(needle) { *xor_counter += 1; return left_expr; } @@ -76,7 +88,7 @@ pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> let mut xor_counter: i32 = 0; let result_expr = recursive_delete_xor_in_expr(expr, needle, &mut xor_counter); - if result_expr == *needle { + if result_expr.normalize_eq(needle) { return needle.clone(); } else if xor_counter % 2 == 0 { if is_left { diff --git a/datafusion/optimizer/src/test/udfs.rs b/datafusion/optimizer/src/test/udfs.rs index ba71b6a04a7a2..72b1a63f6720d 100644 --- a/datafusion/optimizer/src/test/udfs.rs +++ b/datafusion/optimizer/src/test/udfs.rs @@ -19,7 +19,7 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_expr::{ ColumnarValue, Expr, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF, - ScalarUDFImpl, Signature, TypeSignature, + ScalarUDFImpl, Signature, TypeSignature, Volatility, }; /// A configurable test UDF for optimizer tests. @@ -44,7 +44,7 @@ impl PlacementTestUDF { // The actual types don't matter since this UDF is not intended for execution. signature: Signature::new( TypeSignature::OneOf(vec![TypeSignature::Any(1), TypeSignature::Any(2)]), - datafusion_expr::Volatility::Immutable, + Volatility::Immutable, ), placement: ExpressionPlacement::MoveTowardsLeafNodes, id: 0, @@ -64,6 +64,13 @@ impl PlacementTestUDF { self.id = id; self } + + /// Set the volatility of the UDF, so that rules which must not duplicate a + /// volatile computation (e.g. `random()`) can be exercised. + pub fn with_volatility(mut self, volatility: Volatility) -> Self { + self.signature.volatility = volatility; + self + } } impl ScalarUDFImpl for PlacementTestUDF { diff --git a/datafusion/physical-expr-adapter/src/rewrite.rs b/datafusion/physical-expr-adapter/src/rewrite.rs index 7345a587ee6a4..3d31feca96031 100644 --- a/datafusion/physical-expr-adapter/src/rewrite.rs +++ b/datafusion/physical-expr-adapter/src/rewrite.rs @@ -319,6 +319,9 @@ mod tests { assert_eq!(source.name(), "__datafusion_file_row_index"); assert_eq!(source.index(), 2); + // The row index column is at index 2, beyond the user-visible schema. + // When the source column lookup fails, the field name is empty. The + // correct field name would be provided by a parent projection/alias. let input_schema = Schema::new(vec![ Field::new("value", DataType::Int64, true), Field::new("__datafusion_file_row_index", DataType::Int64, false) @@ -328,9 +331,11 @@ mod tests { )])), ]); let return_field = expr.return_field(&input_schema)?; - assert_eq!(return_field.name(), "file_row_index"); + // Field name is empty because column index 2 is beyond the schema + assert_eq!(return_field.name(), ""); assert_eq!(return_field.data_type(), &DataType::Int64); assert!(return_field.is_nullable()); + // Exact target field does not preserve source metadata assert!(return_field.metadata().is_empty()); Ok(()) } diff --git a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs index 7c8cdc3b4c50e..68351a839554a 100644 --- a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs +++ b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs @@ -17,7 +17,9 @@ use arrow::array::{ArrayRef, StringArray}; use criterion::{Criterion, Throughput, criterion_group, criterion_main}; -use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_map::{ + ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType, +}; use std::hint::black_box; use std::sync::Arc; @@ -57,7 +59,13 @@ fn bench_arrow_bytes_map(c: &mut Criterion) { for (name, values) in cases { group.bench_function(name, |b| { b.iter(|| { - let mut map = ArrowBytesMap::::new(OutputType::Utf8); + // The `long_low_cardinality` case is defined by the distinct + // values fitting in the pre-allocated buffer, so this benchmark + // measures the pre-allocating constructor. + let mut map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); let mut next_payload = 0; map.insert_if_new( &values, diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 7543e6b297329..4028520c776d0 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -28,7 +28,7 @@ use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::any::type_name; use std::fmt::Debug; @@ -55,10 +55,21 @@ pub enum OutputType { pub struct ArrowBytesSet(ArrowBytesMap); impl ArrowBytesSet { + /// Creates a set that does not pre-allocate its hash table or value buffer. + /// + /// See [`ArrowBytesMap::new`] for when to prefer this over + /// [`Self::with_capacity`]. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesMap::new(output_type)) } + /// Creates a set with room for `map_capacity` entries. + /// + /// See [`ArrowBytesMap::with_capacity`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self(ArrowBytesMap::with_capacity(output_type, map_capacity)) + } + /// Return the contents of this set and replace it with a new empty /// set with the same output type pub fn take(&mut self) -> Self { @@ -217,8 +228,13 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, - /// Total size of the map in bytes - map_size: usize, + /// Hash table capacity to re-create the map with in [`Self::take`], so a + /// map built with [`Self::with_capacity`] keeps its pre-allocation when it + /// is emptied and reused + initial_map_capacity: usize, + /// Value buffer capacity to re-create the buffer with in [`Self::take`], + /// for the same reason as `initial_map_capacity` + initial_buffer_capacity: usize, /// In progress buffer containing all values buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used @@ -236,20 +252,71 @@ where null: Option<(V, usize)>, } -/// The size, in number of entries, of the initial hash table -const INITIAL_MAP_CAPACITY: usize = 128; -/// The initial size, in bytes, of the string data +/// The size, in number of entries, of the hash table pre-allocated by +/// [`ArrowBytesMap::with_capacity`]. It is a warm up size for maps that go on +/// to hold many values, not a bound on what the map can hold. +pub const INITIAL_MAP_CAPACITY: usize = 128; +/// The size, in bytes, of the string data buffer pre-allocated by +/// [`ArrowBytesMap::with_capacity`] pub const INITIAL_BUFFER_CAPACITY: usize = 8 * 1024; + +/// Appends `value` to a map's value buffer, growing the buffer on a power of +/// two ladder. +/// +/// `Vec` doubles from wherever its first allocation landed, so a buffer started +/// empty by [`ArrowBytesMap::new`] and one started at +/// [`INITIAL_BUFFER_CAPACITY`] by [`ArrowBytesMap::with_capacity`] sit on +/// different ladders, and can hold the same values at capacities differing by +/// up to 2x in either direction depending on the value lengths. Rounding every +/// growth up to a power of two puts both on one ladder, which is what makes a +/// lazily allocated map never larger than a pre-allocated one holding the same +/// values. Growth stays geometric, so appending is still amortized constant +/// time. +fn push_value_bytes(buffer: &mut Vec, value: &[u8]) { + let required = buffer.len() + value.len(); + if required > buffer.capacity() { + let target = required.checked_next_power_of_two().unwrap_or(required); + buffer.reserve_exact(target - buffer.len()); + } + buffer.extend_from_slice(value); +} + impl ArrowBytesMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, { + /// Creates a map that does not pre-allocate its hash table or value buffer. + /// + /// Use this when maps are created in large numbers and most of them stay + /// small, such as the per group `COUNT(DISTINCT)` accumulators that + /// `GroupsAccumulatorAdapter` creates one of per group. There the + /// pre-allocation dwarfs the values the map actually holds. pub fn new(output_type: OutputType) -> Self { + Self::new_inner(output_type, 0, 0) + } + + /// Creates a map whose hash table is pre-allocated for `map_capacity` + /// entries and whose value buffer is pre-allocated with + /// [`INITIAL_BUFFER_CAPACITY`] bytes. + /// + /// Use this for the few long lived maps that are each expected to hold many + /// values, such as the single map backing a `GROUP BY` on one string + /// column. The capacities are preserved across [`Self::take`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self::new_inner(output_type, map_capacity, INITIAL_BUFFER_CAPACITY) + } + + fn new_inner( + output_type: OutputType, + map_capacity: usize, + buffer_capacity: usize, + ) -> Self { Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - map_size: 0, - buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, + initial_buffer_capacity: buffer_capacity, + buffer: Vec::with_capacity(buffer_capacity), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), hashes_buffer: vec![], @@ -260,11 +327,31 @@ where /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.output_type); + let mut new_self = Self::new_inner( + self.output_type, + self.initial_map_capacity, + self.initial_buffer_capacity, + ); swap(self, &mut new_self); new_self } + /// Empties this map and releases every allocation it holds, so + /// [`Self::size`] drops to approximately zero. + /// + /// This is the difference from [`Self::take`]: `take` restores the + /// capacities the map was configured with so the emptied map stays warm for + /// continued use, which is what emitting wants, whereas here the point is + /// to hand the memory back, as before spilling or before a downstream sort. + /// The configured capacities are remembered, so the next [`Self::take`] + /// warms the map up again. + pub fn clear_and_release(&mut self) { + let mut released = Self::new_inner(self.output_type, 0, 0); + released.initial_map_capacity = self.initial_map_capacity; + released.initial_buffer_capacity = self.initial_buffer_capacity; + *self = released; + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -406,7 +493,7 @@ where // Put the small values into buffer and offsets so it appears // the output array, but store the actual bytes inline for // comparison - self.buffer.extend_from_slice(value); + push_value_bytes(&mut self.buffer, value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); let new_header = Entry { @@ -415,11 +502,8 @@ where offset_or_inline: inline, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + self.map + .insert_unique(hash, new_header, |header| header.hash); payload } } @@ -447,7 +531,7 @@ where // appears the output array, and store that offset // so the bytes can be compared if needed let offset = self.buffer.len(); // offset of start for data - self.buffer.extend_from_slice(value); + push_value_bytes(&mut self.buffer, value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); @@ -457,11 +541,8 @@ where offset_or_inline: offset, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + self.map + .insert_unique(hash, new_header, |header| header.hash); payload } }; @@ -486,7 +567,8 @@ where let Self { output_type, map: _, - map_size: _, + initial_map_capacity: _, + initial_buffer_capacity: _, offsets, buffer, random_state: _, @@ -591,7 +673,11 @@ where /// Return the total size, in bytes, of memory used to store the data in /// this set, not including `self` pub fn size(&self) -> usize { - self.map_size + // `HashTable::allocation_size` reports the whole hashbrown allocation, + // which is the entry array plus the control bytes plus the trailing + // group, so it is larger than `capacity() * size_of::>()`. + // It is a constant time layout calculation, not a walk of the table. + self.map.allocation_size() + self.buffer.capacity() * size_of::() + self.offsets.allocated_size() + self.hashes_buffer.allocated_size() @@ -615,7 +701,7 @@ where fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ArrowBytesMap") .field("map", &"") - .field("map_size", &self.map_size) + .field("map_allocation_size", &self.map.allocation_size()) .field("buffer", &self.buffer) .field("random_state", &self.random_state) .field("hashes_buffer", &self.hashes_buffer) @@ -663,6 +749,155 @@ mod tests { use arrow::array::{BinaryArray, LargeBinaryArray, StringArray}; use std::collections::HashMap; + /// A lower bound on the bytes a hashbrown table holding `entries` entries + /// of type `T` must allocate: one entry slot and one control byte each. + /// Derived independently of the production accounting so it can bracket it. + fn min_table_bytes(entries: usize) -> usize { + entries * (size_of::() + 1) + } + + #[test] + fn map_new_does_not_allocate() { + let map = ArrowBytesMap::::new(OutputType::Utf8); + + assert_eq!(map.map.capacity(), 0); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.buffer.capacity(), 0); + // Only the single leading zero offset is allocated. + assert!(map.size() < 128, "expected {} to be tiny", map.size()); + } + + #[test] + fn map_with_capacity_reports_the_real_hash_table_allocation() { + let map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + + // Before this accounting was corrected the map reported its hash table + // as costing zero bytes until the table grew past its pre-allocation. + let table_bytes = map.map.allocation_size(); + let lower_bound = min_table_bytes::>(map.map.capacity()); + assert!( + table_bytes >= lower_bound, + "expected {table_bytes} to be at least {lower_bound}" + ); + assert!( + table_bytes <= 2 * lower_bound + 64, + "expected {table_bytes} to be within a small factor of {lower_bound}" + ); + assert!(map.size() >= table_bytes + INITIAL_BUFFER_CAPACITY); + } + + #[test] + fn take_preserves_the_capacity_the_map_was_built_with() { + let mut preallocated = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + let capacity = preallocated.map.capacity(); + preallocated.take(); + assert_eq!(preallocated.map.capacity(), capacity); + assert_eq!(preallocated.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + lazy.take(); + assert_eq!(lazy.map.capacity(), 0); + assert_eq!(lazy.buffer.capacity(), 0); + } + + #[test] + fn clear_and_release_frees_the_preallocation_that_take_keeps() { + let mut map = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1_000).map(|i| format!("distinct value number {i}")), + )); + map.insert_if_new(&values, |_| (), |_| ()); + + let populated_size = map.size(); + assert!(populated_size > INITIAL_BUFFER_CAPACITY); + + // `take` deliberately keeps the map warm, so it does not release the + // configured capacities. + map.take(); + let taken_size = map.size(); + assert!( + taken_size > INITIAL_BUFFER_CAPACITY, + "expected take to retain the warm up allocations, got {taken_size}" + ); + + map.clear_and_release(); + let released_size = map.size(); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.buffer.capacity(), 0); + assert!( + released_size < 128, + "expected the released map to report approximately zero bytes, got {released_size}" + ); + + // The configured capacities survive, so the map warms back up when it + // is emitted from again. + map.take(); + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY); + } + + #[test] + fn lazy_and_pre_allocated_buffers_grow_on_the_same_ladder() { + // Value lengths chosen so the buffer requirement lands between powers + // of two, which is where the two ladders used to diverge. + for value_len in [9usize, 13, 24, 37] { + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + let mut pre_allocated = ArrowBytesMap::::with_capacity( + OutputType::Utf8, + INITIAL_MAP_CAPACITY, + ); + + for batch in 0..8 { + let values: ArrayRef = + Arc::new(StringArray::from_iter_values((0..1_000).map(|i| { + let value = format!("{}:{i}", batch * 1_000 + i); + format!("{value:value_len$}") + }))); + lazy.insert_if_new(&values, |_| (), |_| ()); + pre_allocated.insert_if_new(&values, |_| (), |_| ()); + + assert_eq!(lazy.buffer.len(), pre_allocated.buffer.len()); + assert_eq!( + lazy.buffer.capacity(), + pre_allocated.buffer.capacity(), + "value length {value_len}, batch {batch}: a buffer that \ + started empty reached {} bytes of capacity against {} for \ + one that started at INITIAL_BUFFER_CAPACITY", + lazy.buffer.capacity(), + pre_allocated.buffer.capacity(), + ); + } + } + } + + #[test] + fn a_lazy_buffer_stays_below_the_pre_allocated_floor_while_it_is_small() { + let mut lazy = ArrowBytesMap::::new(OutputType::Utf8); + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..10).map(|i| format!("distinct value number {i}")), + )); + lazy.insert_if_new(&values, |_| (), |_| ()); + + assert!( + lazy.buffer.capacity() < INITIAL_BUFFER_CAPACITY, + "expected a small lazy buffer to stay under the {INITIAL_BUFFER_CAPACITY} \ + byte pre-allocation, got {}", + lazy.buffer.capacity() + ); + } + #[test] fn string_set_empty() { let mut set = ArrowBytesSet::::new(OutputType::Utf8); diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 0457825decb96..29f4014c5f9a4 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -28,10 +28,9 @@ use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType}; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; -use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_common::utils::proxy::VecAllocExt; use datafusion_common::{Result, exec_err}; use std::fmt::Debug; -use std::mem::size_of; use std::sync::Arc; /// HashSet optimized for storing string or binary values that can produce that @@ -40,10 +39,21 @@ use std::sync::Arc; pub struct ArrowBytesViewSet(ArrowBytesViewMap<()>); impl ArrowBytesViewSet { + /// Creates a set that does not pre-allocate its hash table. + /// + /// See [`ArrowBytesViewMap::new`] for when to prefer this over + /// [`Self::with_capacity`]. pub fn new(output_type: OutputType) -> Self { Self(ArrowBytesViewMap::new(output_type)) } + /// Creates a set with room for `map_capacity` entries. + /// + /// See [`ArrowBytesViewMap::with_capacity`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { + Self(ArrowBytesViewMap::with_capacity(output_type, map_capacity)) + } + /// Inserts each value from `values` into the set pub fn insert(&mut self, values: &ArrayRef) { fn make_payload_fn(_value: Option<&[u8]>) {} @@ -55,9 +65,7 @@ impl ArrowBytesViewSet { /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.0.output_type); - std::mem::swap(self, &mut new_self); - new_self + Self(self.0.take()) } /// Converts this set into a `StringViewArray` or `BinaryViewArray` @@ -129,8 +137,10 @@ where output_type: OutputType, /// Underlying hash set for each distinct value map: hashbrown::hash_table::HashTable>, - /// Total size of the map in bytes - map_size: usize, + /// Hash table capacity to re-create the map with in [`Self::take`], so a + /// map built with [`Self::with_capacity`] keeps its pre-allocation when it + /// is emptied and reused + initial_map_capacity: usize, /// Views for all stored values (in insertion order) views: Vec, @@ -151,21 +161,36 @@ where null: Option<(V, usize)>, } -/// The size, in number of entries, of the initial hash table -const INITIAL_MAP_CAPACITY: usize = 512; +/// The size, in number of entries, of the hash table pre-allocated by +/// [`ArrowBytesViewMap::with_capacity`]. It is a warm up size for maps that go +/// on to hold many values, not a bound on what the map can hold. +pub const INITIAL_MAP_CAPACITY: usize = 512; impl ArrowBytesViewMap where V: Debug + PartialEq + Eq + Clone + Copy + Default, { + /// Creates a map that does not pre-allocate its hash table. + /// + /// Use this when maps are created in large numbers and most of them stay + /// small, such as the per group `COUNT(DISTINCT)` accumulators that + /// `GroupsAccumulatorAdapter` creates one of per group. There the + /// pre-allocation dwarfs the values the map actually holds. pub fn new(output_type: OutputType) -> Self { - let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY); - let map_size = map.capacity() * size_of::>(); + Self::with_capacity(output_type, 0) + } + /// Creates a map whose hash table is pre-allocated for `map_capacity` + /// entries. + /// + /// Use this for the few long lived maps that are each expected to hold many + /// values, such as the single map backing a `GROUP BY` on one string + /// column. The capacity is preserved across [`Self::take`]. + pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self { Self { output_type, - map, - map_size, + map: hashbrown::hash_table::HashTable::with_capacity(map_capacity), + initial_map_capacity: map_capacity, views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -179,11 +204,27 @@ where /// Return the contents of this map and replace it with a new empty map with /// the same output type pub fn take(&mut self) -> Self { - let mut new_self = Self::new(self.output_type); + let mut new_self = + Self::with_capacity(self.output_type, self.initial_map_capacity); std::mem::swap(self, &mut new_self); new_self } + /// Empties this map and releases every allocation it holds, so + /// [`Self::size`] drops to approximately zero. + /// + /// This is the difference from [`Self::take`]: `take` restores the capacity + /// the map was configured with so the emptied map stays warm for continued + /// use, which is what emitting wants, whereas here the point is to hand the + /// memory back, as before spilling or before a downstream sort. The + /// configured capacity is remembered, so the next [`Self::take`] warms the + /// map up again. + pub fn clear_and_release(&mut self) { + let mut released = Self::with_capacity(self.output_type, 0); + released.initial_map_capacity = self.initial_map_capacity; + *self = released; + } + /// Inserts each value from `values` into the map, invoking `payload_fn` for /// each value if *not* already present, deferring the allocation of the /// payload until it is needed. @@ -374,8 +415,7 @@ where payload, }; - self.map - .insert_accounted(new_header, |h| h.hash, &mut self.map_size); + self.map.insert_unique(hash, new_header, |h| h.hash); payload }; observe_payload_fn(payload); @@ -540,13 +580,19 @@ where pub fn size(&self) -> usize { // All fields below own their allocations. Count retained capacity rather // than used length because this value drives memory accounting. + // + // `HashTable::allocation_size` reports the whole hashbrown allocation, + // which is the entry array plus the control bytes plus the trailing + // group, so it is larger than `capacity() * size_of::>()`. It + // is a constant time layout calculation, not a walk of the table. + let map_size = self.map.allocation_size(); let views_size = self.views.allocated_size(); let in_progress_size = self.in_progress.allocated_size(); let completed_size = self.completed.allocated_size() + self.completed.iter().map(Buffer::capacity).sum::(); let nulls_size = self.nulls.allocated_size(); - self.map_size + map_size + views_size + in_progress_size + completed_size @@ -562,7 +608,7 @@ where fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ArrowBytesMap") .field("map", &"") - .field("map_size", &self.map_size) + .field("map_allocation_size", &self.map.allocation_size()) .field("views_len", &self.views.len()) .field("completed_buffers", &self.completed.len()) .field("random_state", &self.random_state) @@ -597,6 +643,7 @@ where mod tests { use arrow::array::{GenericByteViewArray, StringViewArray}; use datafusion_common::HashMap; + use std::mem::size_of; use super::*; @@ -785,11 +832,104 @@ mod tests { assert_eq!(set.len(), 10); } + /// A lower bound on the bytes a hashbrown table holding `entries` entries + /// of type `T` must allocate: one entry slot and one control byte each. + /// Derived independently of the production accounting so it can bracket it. + fn min_table_bytes(entries: usize) -> usize { + entries * (size_of::() + 1) + } + #[test] - fn test_size_counts_initial_hash_table_capacity() { + fn map_new_does_not_allocate() { let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); - assert_eq!(map.size(), map.map.capacity() * size_of::>()); + assert_eq!(map.map.capacity(), 0); + assert_eq!(map.map.allocation_size(), 0); + assert_eq!(map.size(), 0); + } + + #[test] + fn test_size_counts_initial_hash_table_capacity() { + let map = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); + assert_eq!(map.size(), map.map.allocation_size()); + + // Before this accounting was corrected the map reported exactly + // `capacity() * size_of::>()`, which leaves out the control + // bytes and undercounts the real allocation by roughly half. + let lower_bound = min_table_bytes::>(map.map.capacity()); + assert!( + map.size() >= lower_bound, + "expected {} to be at least {lower_bound}", + map.size() + ); + assert!( + map.size() <= 2 * lower_bound + 64, + "expected {} to be within a small factor of {lower_bound}", + map.size() + ); + assert!( + map.size() > map.map.capacity() * size_of::>(), + "expected {} to exceed {}", + map.size(), + map.map.capacity() * size_of::>() + ); + } + + #[test] + fn take_preserves_the_capacity_the_map_was_built_with() { + let mut preallocated = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + let capacity = preallocated.map.capacity(); + preallocated.take(); + assert_eq!(preallocated.map.capacity(), capacity); + + let mut lazy = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); + lazy.take(); + assert_eq!(lazy.map.capacity(), 0); + } + + #[test] + fn clear_and_release_frees_the_preallocation_that_take_keeps() { + let mut map = ArrowBytesViewMap::<()>::with_capacity( + OutputType::Utf8View, + INITIAL_MAP_CAPACITY, + ); + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1_000).map(|i| format!("distinct value number {i}")), + )); + map.insert_if_new(&values, |_| (), |_| ()); + + let warm_size = map.map.allocation_size(); + assert!(warm_size > 0); + + // `take` deliberately keeps the map warm, so it does not release the + // configured capacity. + map.take(); + let taken_size = map.size(); + assert!( + taken_size >= min_table_bytes::>(INITIAL_MAP_CAPACITY), + "expected take to retain the warm up allocation, got {taken_size}" + ); + + map.clear_and_release(); + assert_eq!(map.map.allocation_size(), 0); + let released_size = map.size(); + assert!( + released_size < 128, + "expected the released map to report approximately zero bytes, got {released_size}" + ); + + // The configured capacity survives, so the map warms back up when it is + // emitted from again. + map.take(); + assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY); } #[test] @@ -822,7 +962,7 @@ mod tests { .any(|buffer| buffer.capacity() > buffer.len()) ); - let expected_size = map.map_size + let expected_size = map.map.allocation_size() + map.views.allocated_size() + map.in_progress.allocated_size() + map.completed.allocated_size() @@ -832,7 +972,7 @@ mod tests { assert_eq!(map.size(), expected_size); // Verify the retained-capacity delta independently from the production formula. - let legacy_size = map.map_size + let legacy_size = map.map.allocation_size() + map.views.len() * size_of::() + map.in_progress.capacity() + map.completed.iter().map(Buffer::len).sum::() diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 65ef2a3ceb216..9b0ae96c4f60f 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -51,6 +51,7 @@ proto = [ [dependencies] arrow = { workspace = true } +arrow-schema = { workspace = true } datafusion-common = { workspace = true } datafusion-expr = { workspace = true } datafusion-expr-common = { workspace = true } @@ -104,3 +105,7 @@ name = "string_concat" [package.metadata.cargo-machete] ignored = ["half"] + +[[bench]] +harness = false +name = "equivalence_properties" diff --git a/datafusion/physical-expr/benches/equivalence_properties.rs b/datafusion/physical-expr/benches/equivalence_properties.rs new file mode 100644 index 0000000000000..6993b1056dedc --- /dev/null +++ b/datafusion/physical-expr/benches/equivalence_properties.rs @@ -0,0 +1,169 @@ +// 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. + +//! Benchmarks for the ordering satisfaction checks on [`EquivalenceProperties`]. +//! +//! These are called repeatedly during physical optimization (sort removal, +//! `EnforceSorting`, `EnforceDistribution`, and the requirement checks for +//! windows, joins and aggregates), so their cost shows up directly in planning +//! time. +//! +//! # Scope +//! +//! These measure the satisfaction check itself, not the cost of assembling its +//! arguments. The sort expressions, requirements and orderings are built once, +//! up front. Because the checks take their input by value, each iteration gets a +//! fresh copy from the untimed setup step of `iter_batched`; only the call is +//! timed. Any copying the check does internally is part of what is measured. +//! +//! The benchmarks are parameterized by the number of equivalence classes, since +//! that -- not the schema width, which sits behind an `Arc` -- is what these +//! checks carry around. + +use std::sync::Arc; + +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, + PhysicalSortRequirement, +}; + +fn schema(n_cols: usize) -> SchemaRef { + Arc::new(Schema::new( + (0..n_cols) + .map(|i| Field::new(format!("c{i}"), DataType::Int32, true)) + .collect::>(), + )) +} + +fn col(i: usize) -> Arc { + Arc::new(Column::new(&format!("c{i}"), i)) +} + +fn asc(i: usize) -> PhysicalSortExpr { + PhysicalSortExpr::new(col(i), SortOptions::default()) +} + +/// Properties with three equivalent orderings and `n_classes` equivalence +/// classes, i.e. roughly what a scan feeding a join and a window function looks +/// like. Columns `c0..c7` carry the orderings; the equivalence classes are built +/// from the columns above them. +fn properties(n_classes: usize) -> EquivalenceProperties { + let schema = schema(8 + 2 * n_classes); + let mut props = EquivalenceProperties::new(schema); + props.add_orderings([ + vec![asc(0), asc(1), asc(2), asc(3)], + vec![asc(4), asc(5)], + vec![asc(6)], + ]); + for i in 0..n_classes { + props + .add_equal_conditions(col(8 + 2 * i), col(9 + 2 * i)) + .unwrap(); + } + props +} + +fn bench_ordering_satisfaction(c: &mut Criterion) { + let mut group = c.benchmark_group("equivalence_properties"); + + // Built once; see the "Scope" note at the top of this file. + let one_key = vec![asc(0)]; + // `c7` leads none of the orderings, so this exits on the first key. + let one_key_unsatisfied = vec![asc(7)]; + let four_keys = (0..4).map(asc).collect::>(); + let four_key_ordering = LexOrdering::new(four_keys.clone()).unwrap(); + let one_req = vec![PhysicalSortRequirement::new(col(0), None)]; + let four_reqs = (0..4) + .map(|i| PhysicalSortRequirement::new(col(i), None)) + .collect::>(); + + for n_classes in [2, 8, 32] { + let props = properties(n_classes); + + // A single sort key: the most common shape by far. + group.bench_function( + BenchmarkId::new("ordering_satisfy/1_key", n_classes), + |b| { + b.iter_batched( + || one_key.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("ordering_satisfy/1_key_unsatisfied", n_classes), + |b| { + b.iter_batched( + || one_key_unsatisfied.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + // Four sort keys: exercises the per-key constant registration. + group.bench_function( + BenchmarkId::new("ordering_satisfy/4_keys", n_classes), + |b| { + b.iter_batched( + || four_keys.clone(), + |keys| props.ordering_satisfy(keys).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("ordering_satisfy_requirement/1_key", n_classes), + |b| { + b.iter_batched( + || one_req.clone(), + |reqs| props.ordering_satisfy_requirement(reqs).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("ordering_satisfy_requirement/4_keys", n_classes), + |b| { + b.iter_batched( + || four_reqs.clone(), + |reqs| props.ordering_satisfy_requirement(reqs).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + group.bench_function( + BenchmarkId::new("extract_common_sort_prefix/4_keys", n_classes), + |b| { + b.iter_batched( + || four_key_ordering.clone(), + |ordering| props.extract_common_sort_prefix(ordering).unwrap(), + BatchSize::SmallInput, + ) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_ordering_satisfaction); +criterion_main!(benches); diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 08c05efe0ccc0..c68157fecbd8c 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -640,8 +640,13 @@ impl EquivalenceProperties { return Ok(true); } let schema = self.schema(); - let mut eq_properties = self.clone(); - for element in normal_reqs { + // Registering satisfied keys as constants mutates the state, so it + // needs an owned copy -- but only from the second requirement onwards. + // Single-element requirements (the common case) never pay for the clone. + let last_idx = normal_reqs.len() - 1; + let mut owned = None::; + for (idx, element) in normal_reqs.into_iter().enumerate() { + let eq_properties = owned.as_ref().unwrap_or(self); // Check whether given requirement is satisfied: let ExprProperties { sort_properties, .. @@ -658,10 +663,16 @@ impl EquivalenceProperties { if !satisfy { return Ok(false); } + if idx == last_idx { + // Nothing left to check, so no need to update the state: + break; + } // Treat satisfied keys (and the sub-expressions they pin down) as // constants in subsequent iterations. See // [`Self::add_satisfied_key_constants`] for the rationale. - eq_properties.add_satisfied_key_constants(element.expr)?; + owned + .get_or_insert_with(|| self.clone()) + .add_satisfied_key_constants(element.expr)?; } Ok(true) } @@ -718,8 +729,12 @@ impl EquivalenceProperties { return Ok(full_length); } let schema = self.schema(); - let mut eq_properties = self.clone(); + // Registering satisfied keys as constants mutates the state, so it + // needs an owned copy -- but only from the second sort expression + // onwards. Single-element orderings never pay for the clone. + let mut owned = None::; for (idx, element) in normal_ordering.into_iter().enumerate() { + let eq_properties = owned.as_ref().unwrap_or(self); // Check whether given ordering is satisfied: let ExprProperties { sort_properties, .. @@ -739,10 +754,16 @@ impl EquivalenceProperties { // many we've satisfied so far: return Ok(idx); } + if idx + 1 == full_length { + // Nothing left to check, so no need to update the state: + break; + } // Treat satisfied keys (and the sub-expressions they pin down) as // constants in subsequent iterations. See // [`Self::add_satisfied_key_constants`] for the rationale. - eq_properties.add_satisfied_key_constants(Arc::clone(&element.expr))?; + owned + .get_or_insert_with(|| self.clone()) + .add_satisfied_key_constants(Arc::clone(&element.expr))?; } // All sort expressions are satisfied, return full length: Ok(full_length) diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index e95e9e570cf1b..dd98029c739a8 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -18,10 +18,10 @@ mod literal_lookup_table; use super::{Column, Literal}; -use crate::PhysicalExpr; use crate::expressions::{ CastExpr, LambdaVariable, NegativeExpr, NotExpr, lit, try_cast, }; +use crate::{PhysicalExpr, ScalarFunctionExpr}; use arrow::array::*; use arrow::compute::kernels::zip::zip; use arrow::compute::{ @@ -133,6 +133,7 @@ impl CaseBody { // Determine the set of columns that are used in all the expressions of the case body. // Use an ordered set so lambda variables continue to be positioned after columns let mut used_column_indices = BTreeSet::::new(); + let mut supports_projection = true; let mut collect_column_indices = |expr: &Arc| { expr.apply(|expr| { if let Some(column) = expr.downcast_ref::() { @@ -141,6 +142,12 @@ impl CaseBody { expr.downcast_ref::() { used_column_indices.insert(lambda_variable.index()); + } else if expr.downcast_ref::().is_none() + && expr.downcast_ref::().is_none() + && expr.children().is_empty() + { + // Unknown leaves may read input columns without exposing a Column child. + supports_projection = false; } Ok(TreeNodeRecursion::Continue) }) @@ -216,6 +223,7 @@ impl CaseBody { Ok(ProjectedCaseBody { projection, body: projected_body, + supports_projection, }) } } @@ -247,10 +255,32 @@ impl CaseBody { /// /// The projection vector and the rewritten expression (which only differs from the original in /// column reference indices) are held in a `ProjectedCaseBody`. +/// +/// When `supports_projection` is false, evaluation uses the original body and full batch because +/// the expression may depend on input columns that are not visible in its children. #[derive(Debug, Hash, PartialEq, Eq)] struct ProjectedCaseBody { projection: Vec, body: CaseBody, + /// Whether `body` is safe to evaluate against a projected batch. + /// + /// When false, `projection` and the rewritten `body` must not be used because an + /// expression may read input columns without exposing them through its children. + supports_projection: bool, +} + +impl ProjectedCaseBody { + /// Returns a projection when evaluating the derived body can avoid unused input columns. + fn projection_for(&self, batch: &RecordBatch) -> Option> { + let projection = self + .projection + .iter() + .copied() + .filter(|index| *index < batch.num_columns()) + .collect::>(); + (self.supports_projection && projection.len() < batch.num_columns()) + .then_some(projection) + } } /// The CASE expression is similar to a series of nested if/else and there are two forms that @@ -1049,14 +1079,7 @@ impl CaseExpr { projected: &ProjectedCaseBody, ) -> Result { let return_type = self.data_type(&batch.schema())?; - // projected.projection may include indexes of lambda variables not available on this batch - let projection = projected - .projection - .iter() - .copied() - .filter(|index| *index < batch.num_columns()) - .collect::>(); - if projection.len() < batch.num_columns() { + if let Some(projection) = projected.projection_for(batch) { let projected_batch = batch.project(&projection)?; projected .body @@ -1079,14 +1102,7 @@ impl CaseExpr { projected: &ProjectedCaseBody, ) -> Result { let return_type = self.data_type(&batch.schema())?; - // projected.projection may include indexes of lambda variables not available on this batch - let projection = projected - .projection - .iter() - .copied() - .filter(|index| *index < batch.num_columns()) - .collect::>(); - if projection.len() < batch.num_columns() { + if let Some(projection) = projected.projection_for(batch) { let projected_batch = batch.project(&projection)?; projected .body @@ -1204,23 +1220,14 @@ impl CaseExpr { )?)) } } + } else if let Some(projection) = projected.projection_for(batch) { + // The case expressions do not use all the columns of the input batch. + // Project first to reduce time spent filtering. + let projected_batch = batch.project(&projection)?; + projected.body.expr_or_expr(&projected_batch, when_value) } else { - // projected.projection may include indexes of lambda variables not available on this batch - let projection = projected - .projection - .iter() - .copied() - .filter(|index| *index < batch.num_columns()) - .collect::>(); - if projection.len() < batch.num_columns() { - // The case expressions do not use all the columns of the input batch. - // Project first to reduce time spent filtering. - let projected_batch = batch.project(&projection)?; - projected.body.expr_or_expr(&projected_batch, when_value) - } else { - // All columns are used in the case expressions, so there is no need to project. - self.body.expr_or_expr(batch, when_value) - } + // All columns are used in the case expressions, so there is no need to project. + self.body.expr_or_expr(batch, when_value) } } @@ -1592,6 +1599,7 @@ mod tests { use arrow::datatypes::DataType::Float64; use arrow::datatypes::Field; use datafusion_common::cast::{as_float64_array, as_int32_array}; + use datafusion_common::config::ConfigOptions; use datafusion_common::plan_err; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_expr::type_coercion::binary::type_union_coercion; @@ -1599,6 +1607,55 @@ mod tests { use datafusion_physical_expr_common::physical_expr::fmt_sql; use half::f16; + #[derive(Debug, Hash, PartialEq, Eq)] + struct CustomColumn { + inner: Column, + } + + impl CustomColumn { + fn new(name: &str, index: usize) -> Self { + Self { + inner: Column::new(name, index), + } + } + } + + impl std::fmt::Display for CustomColumn { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + std::fmt::Display::fmt(&self.inner, f) + } + } + + impl PhysicalExpr for CustomColumn { + fn data_type(&self, input_schema: &Schema) -> Result { + self.inner.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result { + self.inner.nullable(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + self.inner.evaluate(batch) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.inner.fmt_sql(f) + } + } + #[test] fn case_with_expr() -> Result<()> { let batch = case_test_batch()?; @@ -1894,6 +1951,127 @@ mod tests { Ok(()) } + #[test] + fn case_without_expr_with_custom_column() -> Result<()> { + let batch = case_test_batch()?; + let schema = batch.schema(); + + let when1 = binary( + Arc::new(CustomColumn::new("a", 0)), + Operator::Eq, + lit("foo"), + &schema, + )?; + let when2 = binary( + Arc::new(CustomColumn::new("a", 0)), + Operator::Eq, + lit("bar"), + &schema, + )?; + let expr = generate_case_when_with_type_coercion( + None, + vec![(when1, lit(123i32)), (when2, lit(456i32))], + None, + schema.as_ref(), + )?; + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = as_int32_array(&result)?; + let expected = Int32Array::from(vec![Some(123), None, None, Some(456)]); + + assert_eq!(&expected, result); + Ok(()) + } + + #[test] + fn case_with_expr_with_custom_column() -> Result<()> { + let batch = case_test_batch_with_extra_columns()?; + let schema = batch.schema(); + + let expr = CaseExpr::try_new( + Some(Arc::new(CustomColumn::new("a", 0))), + vec![(lit("foo"), col("b", &schema)?)], + Some(col("c", &schema)?), + )?; + + match &expr.eval_method { + EvalMethod::WithExpression(projected) => { + assert!(!projected.supports_projection); + } + method => panic!("expected WithExpression, got {method:?}"), + } + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = as_int32_array(&result)?; + let expected = Int32Array::from(vec![Some(1), Some(20), Some(30), Some(40)]); + + assert_eq!(&expected, result); + Ok(()) + } + + #[test] + fn case_expr_or_expr_with_custom_column() -> Result<()> { + let batch = case_test_batch_with_extra_columns()?; + let schema = batch.schema(); + let when = binary( + Arc::new(CustomColumn::new("a", 0)), + Operator::Eq, + lit("foo"), + &schema, + )?; + + let expr = CaseExpr::try_new( + None, + vec![(when, col("b", &schema)?)], + Some(col("c", &schema)?), + )?; + + match &expr.eval_method { + EvalMethod::ExpressionOrExpression(projected) => { + assert!(!projected.supports_projection); + } + method => panic!("expected ExpressionOrExpression, got {method:?}"), + } + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = as_int32_array(&result)?; + let expected = Int32Array::from(vec![Some(1), Some(20), Some(30), Some(40)]); + + assert_eq!(&expected, result); + Ok(()) + } + + #[test] + fn case_expr_or_expr_keeps_nullary_scalar_functions_projectable() -> Result<()> { + let batch = case_test_batch1()?; + let schema = batch.schema(); + let random: Arc = Arc::new(ScalarFunctionExpr::try_new( + datafusion_functions::math::random(), + vec![], + &schema, + Arc::new(ConfigOptions::default()), + )?); + let when = binary(random, Operator::Gt, lit(0.5f64), &schema)?; + + let expr = CaseExpr::try_new( + None, + vec![(when, col("b", &schema)?)], + Some(col("c", &schema)?), + )?; + + match &expr.eval_method { + EvalMethod::ExpressionOrExpression(projected) => { + assert!(projected.supports_projection); + assert_eq!(projected.projection, vec![1, 2]); + } + method => panic!("expected ExpressionOrExpression, got {method:?}"), + } + + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + assert_eq!(result.len(), batch.num_rows()); + Ok(()) + } + #[test] fn case_with_expr_when_null() -> Result<()> { let batch = case_test_batch()?; @@ -1961,6 +2139,22 @@ mod tests { Ok(()) } + fn case_test_batch_with_extra_columns() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ]); + let a = StringArray::from(vec![Some("foo"), Some("baz"), None, Some("bar")]); + let b = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4)]); + let c = Int32Array::from(vec![Some(10), Some(20), Some(30), Some(40)]); + RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(a), Arc::new(b), Arc::new(c)], + ) + .map_err(Into::into) + } + fn case_test_batch1() -> Result { let schema = Schema::new(vec![ Field::new("a", DataType::Int32, true), diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index cb3103d38c52a..565d396dad1b5 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,8 +23,9 @@ use std::sync::Arc; use crate::physical_expr::PhysicalExpr; use arrow::compute::{CastOptions, can_cast_types}; -use arrow::datatypes::{DataType, DataType::*, FieldRef, Schema}; +use arrow::datatypes::{DataType, DataType::*, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::nested_struct::{ @@ -59,8 +61,18 @@ fn can_cast_named_struct_types(source: &DataType, target: &DataType) -> bool { pub struct CastExpr { /// The expression to cast pub expr: Arc, - /// Field metadata describing the desired output after casting + /// The target field. + /// + /// For a type-only cast (see [`CastExpr::new`]) this is a field synthesized + /// from the target data type alone and only its data type is meaningful. + /// For a cast built from an explicit field (see + /// [`CastExpr::new_with_target_field`]) its metadata and nullability are + /// applied to the output field as-is. target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata and + /// nullability describe the output field exactly. + explicit_target: bool, /// Cast options cast_options: CastOptions<'static>, } @@ -68,8 +80,12 @@ pub struct CastExpr { // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for CastExpr { fn eq(&self, other: &Self) -> bool { + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. self.expr.eq(&other.expr) - && self.target_field.eq(&other.target_field) + && self.cast_type().eq(other.cast_type()) + && self.target_metadata().eq(&other.target_metadata()) + && self.target_nullable().eq(&other.target_nullable()) && self.cast_options.eq(&other.cast_options) } } @@ -77,7 +93,17 @@ impl PartialEq for CastExpr { impl Hash for CastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.target_field.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } + self.target_nullable().hash(state); self.cast_options.hash(state); } } @@ -85,39 +111,39 @@ impl Hash for CastExpr { impl CastExpr { /// Create a new `CastExpr` using only a `DataType`. /// - /// This constructor is provided for compatibility with existing call sites - /// that only know the target type. It synthesizes a ``Field`` with the - /// given type (**nullable by default**) and no name metadata. Callers that - /// already have a `FieldRef` (for example, coming from schema inference or a - /// resolved column) should prefer [`CastExpr::new_with_target_field`], which - /// preserves the field's name, nullability, and other metadata. In other - /// words: + /// This constructor creates a type-only cast where metadata and nullability + /// are passed through from the source expression (with extension type keys + /// stripped from metadata). This is the most common use case when you only + /// need to change the data type. /// - /// * use `new()` when only a `DataType` is available and you want the legacy - /// semantics of a type-only cast - /// * use `new_with_target_field()` when you need explicit field - /// metadata/name/nullability preserved + /// For explicit control over the output field's metadata and nullability, + /// use [`CastExpr::new_with_target_field`] or the individual builder methods. pub fn new( expr: Arc, cast_type: DataType, cast_options: Option>, ) -> Self { - Self::new_with_target_field( + Self { expr, - cast_type.into_nullable_field_ref(), - cast_options, - ) + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), + } } /// Create a new `CastExpr` with an explicit target `FieldRef`. /// - /// The provided `target_field` is used verbatim for the expression's - /// return schema, so the field's name, nullability, and other metadata are - /// preserved. This is the preferred constructor when the caller already - /// has field information (for example, during logical-to-physical planning). + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// - The field's nullability is preserved + /// + /// This is the preferred constructor when the caller has explicit field + /// information that should be used exactly (for example, during schema + /// enforcement or adapter layers). /// - /// See [`CastExpr::new`] for the compatibility constructor that only accepts - /// a `DataType`. + /// See [`CastExpr::new`] for type-only casts where source metadata should + /// pass through. pub fn new_with_target_field( expr: Arc, target_field: FieldRef, @@ -126,6 +152,7 @@ impl CastExpr { Self { expr, target_field, + explicit_target: true, cast_options: cast_options.unwrap_or(DEFAULT_CAST_OPTIONS), } } @@ -140,7 +167,30 @@ impl CastExpr { self.target_field.data_type() } - /// Field metadata describing the output column after casting. + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// Explicit nullability for the output field, or `None` to pass through source nullability. + pub fn target_nullable(&self) -> Option { + self.explicit_target + .then(|| self.target_field.is_nullable()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. Note that the returned + /// field may not match what `return_field()` returns when evaluated against + /// a schema, since `return_field()` may incorporate source field information. + /// + /// Prefer [`cast_type()`], [`target_metadata()`], and [`target_nullable()`] + /// for direct access to the individual components. + /// + /// [`cast_type()`]: CastExpr::cast_type + /// [`target_metadata()`]: CastExpr::target_metadata + /// [`target_nullable()`]: CastExpr::target_nullable pub fn target_field(&self) -> &FieldRef { &self.target_field } @@ -150,19 +200,53 @@ impl CastExpr { &self.cast_options } + /// Whether this cast has explicit metadata (vs pass-through from source). + pub fn has_explicit_metadata(&self) -> bool { + self.explicit_target + } + + /// Whether this cast has explicit nullability (vs pass-through from source). + pub fn has_explicit_nullability(&self) -> bool { + self.explicit_target + } + fn resolved_target_field(&self, input_schema: &Schema) -> Result { - if is_default_target_field(&self.target_field) { - self.expr.return_field(input_schema).map(|field| { - Arc::new( - field - .as_ref() - .clone() - .with_data_type(self.cast_type().clone()), + // Try to get the source field for the name. If the target field is + // explicit, we can fall back to an empty name if the source lookup fails + // (e.g., for virtual row-index columns appended at scan time). + let source_result = self.expr.return_field(input_schema); + + if self.explicit_target { + // Metadata and nullability come from the target field verbatim + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new( + name, + self.cast_type().clone(), + self.target_field.is_nullable(), ) - }) - } else { - Ok(Arc::clone(&self.target_field)) + .with_metadata(self.target_field.metadata().clone()), + )); } + + // Type-only cast: pass through the source metadata and nullability, + // stripping extension type keys (the cast is to a plain storage type). + source_result.map(|source_field| { + let mut metadata = source_field.metadata().clone(); + metadata.remove(EXTENSION_TYPE_NAME_KEY); + metadata.remove(EXTENSION_TYPE_METADATA_KEY); + + Arc::new( + source_field + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + .with_metadata(metadata), + ) + }) } /// Check if casting from the specified source type to the target type is a @@ -191,12 +275,6 @@ impl CastExpr { } } -fn is_default_target_field(target_field: &FieldRef) -> bool { - target_field.name().is_empty() - && target_field.is_nullable() - && target_field.metadata().is_empty() -} - pub(crate) fn is_order_preserving_cast_family( source_type: &DataType, target_type: &DataType, @@ -267,11 +345,12 @@ impl PhysicalExpr for CastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(CastExpr::new_with_target_field( - Arc::clone(&children[0]), - Arc::clone(&self.target_field), - Some(self.cast_options.clone()), - ))) + Ok(Arc::new(CastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + cast_options: self.cast_options.clone(), + })) } fn evaluate_bounds(&self, children: &[&Interval]) -> Result { @@ -372,32 +451,61 @@ pub fn cast_with_options( cast_type: DataType, cast_options: Option>, ) -> Result> { - cast_with_target_field( - expr, - input_schema, - cast_type.into_nullable_field_ref(), - cast_options, - ) + let expr_type = expr.data_type(input_schema)?; + + // If the types match, no cast is needed for a type-only cast + if expr_type == cast_type { + return Ok(Arc::clone(&expr)); + } + + let can_build_cast = if requires_nested_struct_cast(&expr_type, &cast_type) { + can_cast_named_struct_types(&expr_type, &cast_type) + } else { + can_cast_types(&expr_type, &cast_type) + }; + + if !can_build_cast { + return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); + } + + Ok(Arc::new(CastExpr::new(expr, cast_type, cast_options))) } /// Return a PhysicalExpression representing `expr` casted to `target_field`, /// preserving any explicit field semantics such as name, nullability, and /// metadata. /// -/// If the input expression already has the same data type, this helper still -/// preserves an explicit `target_field` by constructing a field-aware -/// [`CastExpr`]. Only the default synthesized field created by the legacy -/// type-only API is elided back to the original child expression. +/// If the input expression already has the same data type and the target field +/// has no explicit metadata or nullability constraints, the original expression +/// is returned unchanged. pub fn cast_with_target_field( expr: Arc, input_schema: &Schema, - target_field: FieldRef, + target_field: &FieldRef, cast_options: Option>, ) -> Result> { let expr_type = expr.data_type(input_schema)?; let cast_type = target_field.data_type(); - if expr_type == *cast_type && is_default_target_field(&target_field) { - return Ok(Arc::clone(&expr)); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata or nullability constraints). This is the field created by + // `into_nullable_field_ref()` when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + // For same-type casts, we can skip creating a CastExpr only if: + // 1. The target is type-only (no explicit metadata) + // 2. The source has no extension metadata that needs to be stripped + // Otherwise we need the CastExpr to strip extension metadata from the source. + if expr_type == *cast_type && is_type_only { + let source_field = expr.return_field(input_schema)?; + let has_extension_metadata = source_field + .metadata() + .contains_key(EXTENSION_TYPE_NAME_KEY); + if !has_extension_metadata { + return Ok(Arc::clone(&expr)); + } } let can_build_cast = if requires_nested_struct_cast(&expr_type, cast_type) { @@ -415,11 +523,22 @@ pub fn cast_with_target_field( return not_impl_err!("Unsupported CAST from {expr_type} to {cast_type}"); } - Ok(Arc::new(CastExpr::new_with_target_field( - expr, - target_field, - cast_options, - ))) + // For type-only casts, use CastExpr::new which preserves source metadata/nullability. + // For explicit target fields, use new_with_target_field which applies the target's + // extension metadata and nullability. + if is_type_only { + Ok(Arc::new(CastExpr::new( + expr, + cast_type.clone(), + cast_options, + ))) + } else { + Ok(Arc::new(CastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + cast_options, + ))) + } } /// Return a PhysicalExpression representing `expr` casted to @@ -990,26 +1109,32 @@ mod tests { #[test] fn field_aware_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); for (child_nullable, target_nullable) in [(true, false), (false, true)] { let schema = Schema::new(vec![Field::new("a", Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", Int64, target_nullable) + .with_metadata(metadata.clone()), + ); let expr = CastExpr::new_with_target_field( col("a", &schema)?, - Arc::new( - Field::new("cast_target", Int64, target_nullable) - .with_metadata(metadata.clone()), - ), + Arc::clone(&target_field), None, ); let field = expr.return_field(&schema)?; - assert_eq!(field.name(), "cast_target"); + // Field name comes from source + assert_eq!(field.name(), "a"); assert_eq!(field.data_type(), &Int64); + // Nullability comes from target assert_eq!(field.is_nullable(), target_nullable); + // Target metadata should be preserved exactly assert_eq!( - field.metadata().get("target_meta").map(String::as_str), - Some("1") + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" ); assert_eq!(expr.nullable(&schema)?, child_nullable || target_nullable); } @@ -1017,6 +1142,38 @@ mod tests { Ok(()) } + #[test] + fn target_field_accessor_returns_the_constructed_field() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", Int32, true)]); + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + let target_field = + Arc::new(Field::new("cast_target", Int64, false).with_metadata(metadata)); + + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + // The field is returned verbatim, including its name. + assert_eq!(expr.target_field(), &target_field); + assert_eq!(expr.cast_type(), &Int64); + assert_eq!(expr.target_metadata(), Some(target_field.metadata())); + assert_eq!(expr.target_nullable(), Some(false)); + assert!(expr.has_explicit_metadata()); + assert!(expr.has_explicit_nullability()); + + // A type-only cast reports no explicit target. + let type_only = CastExpr::new(col("a", &schema)?, Int64, None); + assert_eq!(type_only.cast_type(), &Int64); + assert_eq!(type_only.target_metadata(), None); + assert_eq!(type_only.target_nullable(), None); + assert!(!type_only.has_explicit_metadata()); + assert!(!type_only.has_explicit_nullability()); + + Ok(()) + } + #[test] fn type_only_cast_preserves_legacy_field_name_and_nullability() -> Result<()> { let schema = Schema::new(vec![Field::new("a", Int32, false)]); @@ -1170,7 +1327,8 @@ mod tests { let literal = Arc::new(crate::expressions::Literal::new(ScalarValue::Struct( Arc::new(scalar_struct), ))); - let expr = CastExpr::new_with_target_field(literal, Arc::new(target_field), None); + let target_field = Arc::new(target_field); + let expr = CastExpr::new_with_target_field(literal, target_field, None); let batch = RecordBatch::new_empty(schema); let result = expr.evaluate(&batch)?; @@ -1216,6 +1374,93 @@ mod tests { Ok(()) } + #[test] + fn type_only_cast_strips_extension_metadata() -> Result<()> { + // When using type-only cast (new()), extension metadata from source should NOT propagate + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.uuid".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let expr = CastExpr::new(col("a", &schema)?, Utf8, None); + + let field = expr.return_field(&schema)?; + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only cast should strip extension type name from source" + ); + assert_eq!( + field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "Type-only cast should preserve non-extension metadata" + ); + + Ok(()) + } + + #[test] + fn field_aware_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", FixedSizeBinary(16), false).with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", Utf8, true).with_metadata(target_meta)); + let expr = CastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + None, + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware cast should preserve target's non-extension metadata" + ); + + Ok(()) + } + #[test] fn test_check_bigger_cast_precision_loss() { use DataType::*; diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 521b8b87e305c..f2f9285de560a 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -61,7 +61,7 @@ pub use no_op::NoOp; pub use not::{NotExpr, not}; pub(crate) use similar_to_pattern::translate_scalar; pub use similar_to_pattern::{SqlSimilarToPattern, sql_similar_to_regex}; -pub use try_cast::{TryCastExpr, try_cast}; +pub use try_cast::{TryCastExpr, try_cast, try_cast_with_target_field}; pub use unknown_column::UnKnownColumn; pub(crate) use cast::cast_with_target_field; diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index 65b953fd181b7..c054026724fb1 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; use std::fmt; use std::hash::Hash; use std::sync::Arc; @@ -22,40 +23,94 @@ use std::sync::Arc; use crate::PhysicalExpr; use arrow::compute; use arrow::compute::CastOptions; -use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use compute::can_cast_types; +use datafusion_common::datatype::DataTypeExt; use datafusion_common::format::DEFAULT_FORMAT_OPTIONS; use datafusion_common::{Result, not_impl_err}; use datafusion_expr::ColumnarValue; /// TRY_CAST expression casts an expression to a specific data type and returns NULL on invalid cast -#[derive(Debug, Eq)] +#[derive(Debug, Clone, Eq)] pub struct TryCastExpr { /// The expression to cast expr: Arc, - /// The data type to cast to - cast_type: DataType, + /// The target field. + /// + /// For a type-only cast (see [`TryCastExpr::new`]) this is a field + /// synthesized from the target data type alone and only its data type is + /// meaningful. For a cast built from an explicit field (see + /// [`TryCastExpr::new_with_target_field`]) its metadata is applied to the + /// output field as-is. + target_field: FieldRef, + /// Whether `target_field` was supplied by the caller (as opposed to being + /// synthesized from a `DataType`), and therefore whether its metadata + /// describes the output field exactly. + explicit_target: bool, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 impl PartialEq for TryCastExpr { fn eq(&self, other: &Self) -> bool { - self.expr.eq(&other.expr) && self.cast_type == other.cast_type + // Compare the semantically meaningful parts of the target field only: + // the field name never affects the output of this expression. + self.expr.eq(&other.expr) + && self.cast_type() == other.cast_type() + && self.target_metadata() == other.target_metadata() } } impl Hash for TryCastExpr { fn hash(&self, state: &mut H) { self.expr.hash(state); - self.cast_type.hash(state); + self.cast_type().hash(state); + // Hash the metadata by iterating over sorted keys for deterministic ordering + if let Some(metadata) = self.target_metadata() { + let mut entries: Vec<_> = metadata.iter().collect(); + entries.sort_by_key(|(k, _)| *k); + for (k, v) in entries { + k.hash(state); + v.hash(state); + } + } } } impl TryCastExpr { - /// Create a new CastExpr + /// Create a new `TryCastExpr` using only a `DataType`. + /// + /// This constructor creates a type-only cast where metadata is passed through + /// from the source expression (with extension type keys stripped). + /// TRY_CAST results are always nullable since failed casts return NULL. pub fn new(expr: Arc, cast_type: DataType) -> Self { - Self { expr, cast_type } + Self { + expr, + target_field: cast_type.into_nullable_field_ref(), + explicit_target: false, + } + } + + /// Create a new `TryCastExpr` with an explicit target `FieldRef`. + /// + /// The provided `target_field` determines the output characteristics: + /// - The field's data type becomes the cast target type + /// - The field's metadata is used exactly as provided + /// + /// TRY_CAST results are always nullable since failed casts return NULL. + /// + /// See [`TryCastExpr::new`] for type-only casts where source metadata should + /// pass through. + pub fn new_with_target_field( + expr: Arc, + target_field: FieldRef, + ) -> Self { + Self { + expr, + target_field, + explicit_target: true, + } } /// The expression to cast @@ -65,19 +120,33 @@ impl TryCastExpr { /// The data type to cast to pub fn cast_type(&self) -> &DataType { - &self.cast_type + self.target_field.data_type() + } + + /// Explicit metadata for the output field, or `None` to pass through source metadata. + pub fn target_metadata(&self) -> Option<&HashMap> { + self.explicit_target.then(|| self.target_field.metadata()) + } + + /// The target field this cast was constructed with. + /// + /// For a type-only cast this is a field synthesized from the target data + /// type alone; only its data type is meaningful. TRY_CAST results are + /// always nullable regardless of the target field's nullability. + pub fn target_field(&self) -> &FieldRef { + &self.target_field } } impl fmt::Display for TryCastExpr { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type) + write!(f, "TRY_CAST({} AS {})", self.expr, self.cast_type()) } } impl PhysicalExpr for TryCastExpr { fn data_type(&self, _input_schema: &Schema) -> Result { - Ok(self.cast_type.clone()) + Ok(self.cast_type().clone()) } fn nullable(&self, _input_schema: &Schema) -> Result { @@ -90,14 +159,41 @@ impl PhysicalExpr for TryCastExpr { safe: true, format_options: DEFAULT_FORMAT_OPTIONS, }; - value.cast_to(&self.cast_type, Some(&options)) + value.cast_to(self.cast_type(), Some(&options)) } fn return_field(&self, input_schema: &Schema) -> Result { - self.expr - .return_field(input_schema) - .map(|f| f.as_ref().clone().with_data_type(self.cast_type.clone())) - .map(Arc::new) + // If metadata is explicit, we can build the field without source + // (though we still try to get source for the name) + let source_result = self.expr.return_field(input_schema); + + if let Some(metadata) = self.target_metadata() { + // Explicit metadata: use it exactly, TRY_CAST is always nullable + let name = source_result + .as_ref() + .map(|f| f.name().to_string()) + .unwrap_or_default(); + return Ok(Arc::new( + Field::new(name, self.cast_type().clone(), true) + .with_metadata(metadata.clone()), + )); + } + + // Pass-through metadata from source (stripping extension keys) + source_result.map(|source_field| { + let mut metadata = source_field.metadata().clone(); + metadata.remove(EXTENSION_TYPE_NAME_KEY); + metadata.remove(EXTENSION_TYPE_METADATA_KEY); + + Arc::new( + source_field + .as_ref() + .clone() + .with_data_type(self.cast_type().clone()) + .with_nullable(true) // TRY_CAST is always nullable + .with_metadata(metadata), + ) + }) } fn children(&self) -> Vec<&Arc> { @@ -108,16 +204,17 @@ impl PhysicalExpr for TryCastExpr { self: Arc, children: Vec>, ) -> Result> { - Ok(Arc::new(TryCastExpr::new( - Arc::clone(&children[0]), - self.cast_type.clone(), - ))) + Ok(Arc::new(TryCastExpr { + expr: Arc::clone(&children[0]), + target_field: Arc::clone(&self.target_field), + explicit_target: self.explicit_target, + })) } fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "TRY_CAST(")?; self.expr.fmt_sql(f)?; - write!(f, " AS {:?})", self.cast_type) + write!(f, " AS {:?})", self.cast_type()) } #[cfg(feature = "proto")] @@ -190,6 +287,60 @@ pub fn try_cast( } } +/// Return a PhysicalExpression representing `expr` casted to `target_field`, +/// preserving any explicit field semantics such as metadata. +/// +/// TRY_CAST results are always nullable since failed casts return NULL. +/// +/// If the input expression already has the same data type, the target field +/// has no explicit metadata constraints, and the source has no extension +/// metadata to strip, the original expression is returned unchanged. +pub fn try_cast_with_target_field( + expr: Arc, + input_schema: &Schema, + target_field: &FieldRef, +) -> Result> { + let expr_type = expr.data_type(input_schema)?; + let cast_type = target_field.data_type(); + + // Check if this is a "default" target field (type-only cast with no explicit + // metadata constraints). This is the field created by `into_nullable_field_ref()` + // when only a DataType is known. + let is_type_only = target_field.name().is_empty() + && target_field.is_nullable() + && target_field.metadata().is_empty(); + + // For same-type casts, we can skip creating a TryCastExpr only if: + // 1. The target is type-only (no explicit metadata) + // 2. The source has no extension metadata that needs to be stripped + // Otherwise we need the TryCastExpr to strip extension metadata from the source. + if expr_type == *cast_type && is_type_only { + let source_field = expr.return_field(input_schema)?; + let has_extension_metadata = source_field + .metadata() + .contains_key(EXTENSION_TYPE_NAME_KEY); + if !has_extension_metadata { + return Ok(Arc::clone(&expr)); + } + } + + if !can_cast_types(&expr_type, cast_type) { + return not_impl_err!("Unsupported TRY_CAST from {expr_type} to {cast_type}"); + } + + // For type-only casts, use TryCastExpr::new which preserves source metadata. + // For explicit target fields, use new_with_target_field which applies the target's + // metadata exactly. + if is_type_only { + Ok(Arc::new(TryCastExpr::new(expr, cast_type.clone()))) + } else { + Ok(Arc::new(TryCastExpr::new_with_target_field( + expr, + Arc::clone(target_field), + ))) + } +} + #[cfg(test)] mod tests { use super::*; @@ -642,6 +793,167 @@ mod tests { Ok(()) } + + #[test] + fn field_aware_try_cast_uses_exact_target_metadata() -> Result<()> { + // When using field-aware cast, target's metadata should be used exactly + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.type".to_string(), + ), + ("source_key".to_string(), "source_value".to_string()), + ]); + let target_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "target.type".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "target_ext_meta".to_string(), + ), + ("target_key".to_string(), "target_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(16), false) + .with_metadata(source_meta), + ]); + + let target_field = + Arc::new(Field::new("b", DataType::Utf8, true).with_metadata(target_meta)); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"target.type".to_string()), + "Field-aware try_cast should use target's extension type name" + ); + assert_eq!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"target_ext_meta".to_string()), + "Field-aware try_cast should use target's extension type metadata" + ); + assert!( + field.metadata().get("source_key").is_none(), + "Field-aware try_cast should NOT preserve source metadata" + ); + assert_eq!( + field.metadata().get("target_key"), + Some(&"target_value".to_string()), + "Field-aware try_cast should preserve target's non-extension metadata" + ); + // TRY_CAST is always nullable + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn field_aware_try_cast_preserves_target_field_semantics() -> Result<()> { + // Target field metadata should be preserved exactly (no merging with source). + // TRY_CAST is always nullable regardless of target field's nullability. + let metadata = HashMap::from([("target_meta".to_string(), "1".to_string())]); + + for child_nullable in [true, false] { + let schema = + Schema::new(vec![Field::new("a", DataType::Int32, child_nullable)]); + let target_field = Arc::new( + Field::new("cast_target", DataType::Int64, false) // target says non-nullable + .with_metadata(metadata.clone()), + ); + let expr = TryCastExpr::new_with_target_field( + col("a", &schema)?, + Arc::clone(&target_field), + ); + + let field = expr.return_field(&schema)?; + // Field name comes from source + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + // TRY_CAST is ALWAYS nullable (ignores target field's nullability) + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + // Target metadata should be preserved exactly + assert_eq!( + field.metadata().get("target_meta"), + Some(&"1".to_string()), + "Target metadata should be preserved exactly" + ); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + } + + Ok(()) + } + + #[test] + fn type_only_try_cast_strips_extension_keys() -> Result<()> { + // Type-only cast should strip extension keys but preserve other source metadata + let source_meta = HashMap::from([ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "source.extension".to_string(), + ), + ( + EXTENSION_TYPE_METADATA_KEY.to_string(), + "ext_meta".to_string(), + ), + ("custom_key".to_string(), "custom_value".to_string()), + ]); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false).with_metadata(source_meta), + ]); + + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + let field = expr.return_field(&schema)?; + + // Extension keys should be stripped + assert!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "Type-only try_cast should strip extension type name" + ); + assert!( + field.metadata().get(EXTENSION_TYPE_METADATA_KEY).is_none(), + "Type-only try_cast should strip extension type metadata" + ); + // Non-extension metadata should pass through + assert_eq!( + field.metadata().get("custom_key"), + Some(&"custom_value".to_string()), + "Type-only try_cast should preserve non-extension metadata" + ); + // Field name preserved, type changed, always nullable + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable()); + + Ok(()) + } + + #[test] + fn type_only_try_cast_is_always_nullable() -> Result<()> { + // TRY_CAST is always nullable even when source is non-nullable + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let expr = TryCastExpr::new(col("a", &schema)?, DataType::Int64); + + let field = expr.return_field(&schema)?; + + assert_eq!(field.name(), "a"); + assert_eq!(field.data_type(), &DataType::Int64); + assert!(field.is_nullable(), "TRY_CAST should always be nullable"); + assert!( + expr.nullable(&schema)?, + "TRY_CAST should always be nullable" + ); + + Ok(()) + } } #[cfg(all(test, feature = "proto"))] diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 9a0bdc33da8e9..8c5347db742bb 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -27,7 +27,7 @@ use crate::{ use arrow::datatypes::Schema; use datafusion_common::config::ConfigOptions; use datafusion_common::datatype::FieldExt; -use datafusion_common::metadata::{FieldMetadata, format_type_and_metadata}; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{ DFSchema, Result, ScalarValue, TableReference, ToDFSchema, exec_err, internal_datafusion_err, not_impl_err, plan_datafusion_err, plan_err, @@ -387,23 +387,11 @@ pub fn create_physical_expr( Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, - Arc::clone(field), + field, None, ), Expr::TryCast(TryCast { expr, field }) => { - if !field.metadata().is_empty() { - let (_, src_field) = expr.to_field(input_dfschema)?; - return plan_err!( - "TryCast from {} to {} is not supported", - format_type_and_metadata( - src_field.data_type(), - Some(src_field.metadata()), - ), - format_type_and_metadata(field.data_type(), Some(field.metadata())) - ); - } - - expressions::try_cast( + expressions::try_cast_with_target_field( create_physical_expr( expr, input_dfschema, @@ -411,7 +399,7 @@ pub fn create_physical_expr( planning_ctx, )?, input_schema, - field.data_type().clone(), + field, ) } Expr::Not(expr) => expressions::not(create_physical_expr( @@ -743,6 +731,7 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { mod tests { use arrow::array::{ArrayRef, BooleanArray, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field}; + use arrow_schema::extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; use datafusion_common::HashMap; use datafusion_expr::physical_planning_context::{ ScalarSubqueryResults, SubqueryIndex, @@ -771,6 +760,14 @@ mod tests { .expect("planner should lower logical CAST to CastExpr") } + fn as_planner_try_cast( + physical: &Arc, + ) -> &expressions::TryCastExpr { + physical + .downcast_ref::() + .expect("planner should lower logical TRY_CAST to TryCastExpr") + } + #[test] fn test_create_physical_expr_scalar_input_output() -> Result<()> { let expr = col("letter").eq(lit("A")); @@ -833,9 +830,21 @@ mod tests { #[test] fn test_cast_lowering_preserves_target_field_metadata() -> Result<()> { let schema = test_cast_schema(); + + // Target field with both extension metadata and custom metadata. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( - Field::new("cast_target", DataType::Int64, true) - .with_metadata([("target_meta".to_string(), "1".to_string())].into()), + Field::new("cast_target", DataType::Int64, true).with_metadata( + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.json".to_string(), + ), + (EXTENSION_TYPE_METADATA_KEY.to_string(), "{}".to_string()), + ("custom_target_meta".to_string(), "custom_value".to_string()), + ] + .into(), + ), ); let cast_expr = Expr::Cast(Cast::new_from_field( Box::new(col("a")), @@ -845,8 +854,35 @@ mod tests { let physical = lower_cast_expr(&cast_expr, &schema)?; let cast = as_planner_cast(&physical); - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); + // The CastExpr stores the target type and all target metadata + assert_eq!(cast.cast_type(), &DataType::Int64); + let target_metadata = cast.target_metadata().expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + assert_eq!(cast.target_nullable(), Some(true)); + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.json".to_string()) + ); + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_METADATA_KEY), + Some(&"{}".to_string()) + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_target_meta"), + Some(&"custom_value".to_string()), + "All target metadata should propagate with exact semantics" + ); assert!(physical.nullable(&schema)?); Ok(()) @@ -872,22 +908,85 @@ mod tests { #[test] fn test_cast_lowering_preserves_same_type_field_semantics() -> Result<()> { let schema = test_cast_schema(); + + // Same-type cast with extension metadata on target. + // With exact target metadata semantics, all target metadata should propagate. let target_field = Arc::new( Field::new("same_type_cast", DataType::Int32, true).with_metadata( - [("target_meta".to_string(), "same-type".to_string())].into(), + [ + ( + EXTENSION_TYPE_NAME_KEY.to_string(), + "arrow.opaque".to_string(), + ), + ("custom_meta".to_string(), "custom_value".to_string()), + ] + .into(), ), ); - let cast_expr = Expr::Cast(Cast::new_from_field( - Box::new(col("a")), - Arc::clone(&target_field), - )); - let physical = lower_cast_expr(&cast_expr, &schema)?; - let cast = as_planner_cast(&physical); + for use_try_cast in [false, true] { + // For error labelling + let cast_name = if use_try_cast { "TRY_CAST" } else { "CAST" }; - assert_eq!(cast.target_field(), &target_field); - assert_eq!(physical.return_field(&schema)?, target_field); - assert!(physical.nullable(&schema)?); + let cast_expr = if use_try_cast { + Expr::TryCast(TryCast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + } else { + Expr::Cast(Cast::new_from_field( + Box::new(col("a")), + Arc::clone(&target_field), + )) + }; + + let physical = lower_cast_expr(&cast_expr, &schema)?; + + // Extract common fields - both CastExpr and TryCastExpr have these + let (cast_type, target_metadata, target_nullable) = if use_try_cast { + let cast = as_planner_try_cast(&physical); + (cast.cast_type(), cast.target_metadata(), None) + } else { + let cast = as_planner_cast(&physical); + ( + cast.cast_type(), + cast.target_metadata(), + cast.target_nullable(), + ) + }; + + // Verify the physical expression stores correct metadata (same for both) + assert_eq!(cast_type, &DataType::Int32, "{cast_name}: cast_type"); + let target_metadata = target_metadata.expect("should have metadata"); + assert_eq!( + target_metadata.get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: extension type name" + ); + + // Only CastExpr tracks target_nullable (TryCast is always nullable) + if !use_try_cast { + assert_eq!(target_nullable, Some(true), "{cast_name}: target_nullable"); + } + + // return_field should have all target metadata (exact semantics) + let returned = physical.return_field(&schema)?; + assert_eq!( + returned.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"arrow.opaque".to_string()), + "{cast_name}: return_field extension type name" + ); + // All target metadata should propagate with exact semantics + assert_eq!( + returned.metadata().get("custom_meta"), + Some(&"custom_value".to_string()), + "{cast_name}: All target metadata should propagate with exact semantics" + ); + assert!( + physical.nullable(&schema)?, + "{cast_name}: should be nullable" + ); + } Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs index 34ec36be31d2e..9f7b4b4e91cba 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs @@ -22,7 +22,9 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; use datafusion_common::Result; use datafusion_expr::{EmitTo, GroupSelection}; -use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_map::{ + ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType, +}; /// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values /// @@ -38,7 +40,9 @@ pub struct GroupValuesBytes { impl GroupValuesBytes { pub fn new(output_type: OutputType) -> Self { Self { - map: ArrowBytesMap::new(output_type), + // One map holds every group value for the whole query, so it is + // worth pre-allocating the hash table and the value buffer. + map: ArrowBytesMap::with_capacity(output_type, INITIAL_MAP_CAPACITY), num_groups: 0, } } @@ -133,8 +137,67 @@ impl GroupValues for GroupValuesBytes { } fn clear_shrink(&mut self, _num_rows: usize) { - // in theory we could potentially avoid this reallocation and clear the - // contents of the maps, but for now we just reset the map from the beginning - self.map.take(); + // Callers use this to hand memory back before spilling or sorting, so + // release the map's allocations rather than restoring the warm up + // capacities that `take` keeps for the emit path. + self.map.clear_and_release(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::StringArray; + use datafusion_physical_expr_common::binary_map::INITIAL_BUFFER_CAPACITY; + + /// `clear_shrink` is how the aggregate stream hands memory back before it + /// spills and before the spilled batch is sorted, so the memory it releases + /// has to actually show up in the size it reports afterwards. + #[test] + fn clear_shrink_releases_the_reported_memory() { + let mut group_values = GroupValuesBytes::::new(OutputType::Utf8); + let empty = size_of::>(); + + // The map is pre-allocated at construction, so it is already well above + // its own struct size before a single row is interned. + let warm_size = group_values.size(); + assert!( + warm_size > empty + INITIAL_BUFFER_CAPACITY, + "expected the pre-allocated map to report more than {} bytes, got {warm_size}", + empty + INITIAL_BUFFER_CAPACITY + ); + + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1_000).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + let populated_size = group_values.size(); + assert!(populated_size > warm_size); + + group_values.clear_shrink(0); + + // Everything the map held is gone: what remains is the struct itself + // plus the single leading zero offset. + let released_size = group_values.size(); + assert!( + released_size < empty + 128, + "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}" + ); + assert!( + released_size * 10 < populated_size, + "expected {released_size} to be far below {populated_size}" + ); + + // The map still works, and warms back up on the next emit. + group_values.intern(&[values], &mut groups).unwrap(); + assert!(group_values.size() > released_size); + group_values.emit(EmitTo::All).unwrap(); + assert!(group_values.size() > empty + INITIAL_BUFFER_CAPACITY); } } diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs index 997a7ce166a71..23ea4e7ed3f88 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs @@ -19,7 +19,9 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{Array, ArrayRef}; use datafusion_expr::{EmitTo, GroupSelection}; use datafusion_physical_expr::binary_map::OutputType; -use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; +use datafusion_physical_expr_common::binary_view_map::{ + ArrowBytesViewMap, INITIAL_MAP_CAPACITY, +}; use std::mem::size_of; /// A [`GroupValues`] storing single column of Utf8View/BinaryView values @@ -36,7 +38,9 @@ pub struct GroupValuesBytesView { impl GroupValuesBytesView { pub fn new(output_type: OutputType) -> Self { Self { - map: ArrowBytesViewMap::new(output_type), + // One map holds every group value for the whole query, so it is + // worth pre-allocating the hash table. + map: ArrowBytesViewMap::with_capacity(output_type, INITIAL_MAP_CAPACITY), num_groups: 0, } } @@ -135,8 +139,66 @@ impl GroupValues for GroupValuesBytesView { } fn clear_shrink(&mut self, _num_rows: usize) { - // in theory we could potentially avoid this reallocation and clear the - // contents of the maps, but for now we just reset the map from the beginning - self.map.take(); + // Callers use this to hand memory back before spilling or sorting, so + // release the map's allocations rather than restoring the warm up + // capacity that `take` keeps for the emit path. + self.map.clear_and_release(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::StringViewArray; + + /// `clear_shrink` is how the aggregate stream hands memory back before it + /// spills and before the spilled batch is sorted, so the memory it releases + /// has to actually show up in the size it reports afterwards. + #[test] + fn clear_shrink_releases_the_reported_memory() { + let mut group_values = GroupValuesBytesView::new(OutputType::Utf8View); + let empty = size_of::(); + + // The hash table is pre-allocated at construction, so the map is + // already well above its own struct size before a single row is + // interned. + let warm_size = group_values.size(); + assert!( + warm_size > empty + INITIAL_MAP_CAPACITY, + "expected the pre-allocated map to report more than {} bytes, got {warm_size}", + empty + INITIAL_MAP_CAPACITY + ); + + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1_000).map(|i| format!("group value number {i}")), + )); + let mut groups = vec![]; + group_values + .intern(&[Arc::clone(&values)], &mut groups) + .unwrap(); + let populated_size = group_values.size(); + assert!(populated_size > warm_size); + + group_values.clear_shrink(0); + + // Everything the map held is gone: what remains is the struct itself. + let released_size = group_values.size(); + assert!( + released_size < empty + 128, + "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}" + ); + assert!( + released_size * 10 < populated_size, + "expected {released_size} to be far below {populated_size}" + ); + + // The map still works, and warms back up on the next emit. + group_values.intern(&[values], &mut groups).unwrap(); + assert!(group_values.size() > released_size); + group_values.emit(EmitTo::All).unwrap(); + assert!(group_values.size() > empty + INITIAL_MAP_CAPACITY); } } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 340bf5cfc12d6..edf084ad328bc 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -32,9 +32,12 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; -use datafusion_execution::TaskContext; +use datafusion_common::{ + DataFusionError, Result, assert_ne_or_internal_err, internal_datafusion_err, + internal_err, +}; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -53,7 +56,7 @@ use crate::metrics::{ use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; -use crate::stream::EmptyRecordBatchStream; +use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; /// Hash aggregation is implemented in two stages: partial and final. This @@ -249,47 +252,15 @@ pub(crate) struct FinalHashAggregateStream { /// See comments for the same variable in [`PartialHashAggregateStream`]. group_values_soft_limit: Option, - /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// The hash table owns the lower-level /// state for emitting output batches. - state: Option, -} - -/// States for final hash aggregation processing. -// The typestate pattern is used in case the inner logic becomes more complex in -// the future. -enum FinalHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - /// `None` if spilling is not supported by the configured `DiskManager`. - spill_context: Option>, - }, - Spilling { - hash_table: AggregateHashTable, - spill_context: Box, - }, - ProducingOutput { - hash_table: AggregateHashTable, - }, - PreparingMergeInput { - hash_table: AggregateHashTable, - spill_context: Box, - }, - MergingSpills { - stream: SendableRecordBatchStream, - }, - Done, - /// Sentinel state to use when returning error from any other states, because: - /// - It explicitly releases state-owned resources immediately - /// - More defensive against accidentally resuming execution after error - Error, + /// + /// This will be None when creating the stream + hash_table: Option>, + /// `None` if spilling is not supported by the configured `DiskManager`. + spill_context: Option>, } -type FinalHashAggregatePoll = Poll>>; -type FinalHashAggregateStateTransition = ControlFlow< - (FinalHashAggregatePoll, FinalHashAggregateState), - FinalHashAggregateState, ->; - impl FinalSpillContext { fn new( agg: &AggregateExec, @@ -1057,29 +1028,50 @@ impl FinalHashAggregateStream { baseline_metrics, reservation, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), - state: Some(FinalHashAggregateState::ReadingInput { - hash_table, - spill_context, - }), + hash_table: Some(hash_table), + spill_context, }) } - fn close_input(&mut self) { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); + + Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())) } - fn break_with_err(error: DataFusionError) -> FinalHashAggregateStateTransition { - ControlFlow::Break(( - Poll::Ready(Some(Err(error))), - FinalHashAggregateState::Error, - )) + /// Entry point for the final hash aggregate flow + /// + /// See comments in [`FinalHashAggregateStream`] for high-level ideas. + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|emitter| async move { + let mut hash_table = self + .hash_table + .take() + .expect("hash_table should not be None"); + + let mut spill_context = self.spill_context.take(); + + self.consume_input(&mut hash_table, &mut spill_context) + .await?; + self.close_input(); + + match spill_context.filter(|s| s.has_spills()) { + // - If spilled before, perform merging spill runs + Some(spill_context) => { + self.produce_output_from_spills(hash_table, spill_context, emitter) + .await? + } + // Either all the input fit in memory or hit soft group limit with no spilling + None => self.produce_output_from_memory(hash_table, emitter).await?, + } + + Ok(()) + }) } - fn break_with_internal_err( - message: impl std::fmt::Display, - ) -> FinalHashAggregateStateTransition { - Self::break_with_err(internal_datafusion_err!("{message}")) + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } /// See comments in [`Self::group_values_soft_limit`] for details. @@ -1088,14 +1080,6 @@ impl FinalHashAggregateStream { .is_some_and(|limit| limit <= hash_table.building_group_count()) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { - self.close_input(); - hash_table.start_output() - } - /// Reserve memory for the current aggregate table. fn reservation_size_for_table( hash_table: &AggregateHashTable, @@ -1116,454 +1100,174 @@ impl FinalHashAggregateStream { } } - /// Handle ReadingInput state - aggregate partial state batches into the hash table. - /// - /// See comments at `poll_next()` for details. + /// Read input stream, if no memory, then spill and continue reading - aggregate partial state batches into the hash table. /// - /// Returns the next operator state with control flow decision. - fn handle_reading_input( + /// Spilling: The table cannot reserve enough memory. + /// Move all current states into one fully group-key-sorted spill run. + async fn consume_input( &mut self, - cx: &mut Context<'_>, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::ReadingInput { - mut hash_table, - spill_context, - } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected ReadingInput state", - ); - }; + hash_table: &mut AggregateHashTable, + spill_context: &mut Option>, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - FinalHashAggregateState::ReadingInput { - hash_table, - spill_context, - }, - )), - Poll::Ready(Some(Ok(batch))) => { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.aggregate_batch(&batch); - timer.done(); + while let Some(batch) = self.input.next().await.transpose()? { + let _timer = elapsed_compute.timer(); + hash_table.aggregate_batch(&batch)?; + + // Soft group limits are usually small and rarely coincide with + // spilling. Once spilling has occurred, skip this optimization to + // make the internal logic simpler. + let spilled = spill_context + .as_ref() + .is_some_and(|context| context.has_spills()); + if self.hit_soft_group_limit(hash_table) && !spilled { + break; + } - if let Err(e) = result { - return Self::break_with_err(e); - } + // Check memory reservation, and potentially spill. + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + hash_table, + spill_context.as_deref(), + )); + + match resize_result { + Ok(()) => {} + + // The table cannot reserve enough memory. + // Move all current states into one fully group-key-sorted spill run. + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + // OOM and don't support spilling from configuration + let spill_context = spill_context.as_mut().ok_or_else(|| e.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + ))?; + + // Sanity check: impossible to OOM when there is no group aggregated. + assert_ne_or_internal_err!( + hash_table.building_group_count(), + 0, + "Final hash aggregate ran out of memory with no aggregated groups" + ); - // Soft group limits are usually small and rarely coincide with - // spilling. Once spilling has occurred, skip this optimization to - // make the internal logic simpler. - let spilled = spill_context - .as_ref() - .is_some_and(|context| context.has_spills()); - if self.hit_soft_group_limit(&hash_table) && !spilled { - let timer = elapsed_compute.timer(); - let result = self.start_output(&mut hash_table); - timer.done(); + // Sorts and spills one complete in-memory state run - return match result { - Ok(()) => ControlFlow::Continue( - FinalHashAggregateState::ProducingOutput { hash_table }, - ), - Err(e) => Self::break_with_err(e), - }; - } + // Go to the next state to perform spilling the aggregated + // groups so far. + let result = spill_context.spill_table(hash_table); - // Check memory reservation, and potentially spill. - let timer = elapsed_compute.timer(); - let resize_result = + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. self.reservation - .try_resize(Self::reservation_size_for_table( - &hash_table, - spill_context.as_deref(), - )); - timer.done(); - match resize_result { - Ok(()) => {} - Err(e @ DataFusionError::ResourcesExhausted(_)) => { - // OOM and don't support spilling from configuration - let Some(spill_context) = spill_context else { - return Self::break_with_err(e.context( - "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", - )); - }; - // Sanity check: impossible to OOM when there is no group aggregated. - if hash_table.building_group_count() == 0 { - return Self::break_with_internal_err( - "Final hash aggregate ran out of memory with no aggregated groups", - ); - } - // Go to the next state to perform spilling the aggregated - // groups so far. - return ControlFlow::Continue( - FinalHashAggregateState::Spilling { - hash_table, - spill_context, - }, - ); - } - Err(e) => return Self::break_with_err(e), - } + .try_resize(hash_table.memory_size()) + .map_err(|e| { + e.context( + "Decreasing allocation after spilling should succeed", + ) + })?; - ControlFlow::Continue(FinalHashAggregateState::ReadingInput { - hash_table, - spill_context, - }) - } - Poll::Ready(Some(Err(e))) => Self::break_with_err(e), - // Input done, move to next state: - // - If spilled before, perform merging spill runs - // - If not spilled, start producing outputs - Poll::Ready(None) => { - self.close_input(); - match spill_context { - Some(spill_context) if spill_context.has_spills() => { - ControlFlow::Continue( - FinalHashAggregateState::PreparingMergeInput { - hash_table, - spill_context, - }, - ) - } - _ => { - let elapsed_compute = - self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.start_output(); - timer.done(); - - match result { - Ok(()) => ControlFlow::Continue( - FinalHashAggregateState::ProducingOutput { hash_table }, - ), - Err(e) => Self::break_with_err(e), - } - } + result?; + + // One sorted run was written; resume reading the original input. } + Err(e) => return Err(e), } } + + Ok(()) } - /// Sorts and spills one complete in-memory state run, then resumes input. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_spilling( + /// Produce output from spills + /// 1. Spill in progress in-memory hash table + /// 2. Switch to ordered final stream + /// 3. passthrough stream output + async fn produce_output_from_spills( &mut self, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::Spilling { - mut hash_table, - mut spill_context, - } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected Spilling state", - ); - }; - - // Sanity check: it is impossible to OOM when the table is empty. - if hash_table.building_group_count() == 0 { - return Self::break_with_internal_err( - "Final hash aggregation entered Spilling with an empty table", - ); - } - + mut hash_table: AggregateHashTable, + mut spill_context: Box, + mut emitter: TryEmitter, + ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let mut result = spill_context.spill_table(&mut hash_table); - // Spilling shrinks the aggregate table and releases its accumulated - // memory. Update the reservation accordingly. - if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) { - result = - Err(e.context("Decreasing allocation after spilling should succeed")); - } + // Input was exhausted after spilling. Spill the last in-memory run + spill_context.spill_table(&mut hash_table)?; + + // Construct the ordered input used to merge all spill files. + let mut output_stream = + self.switch_to_ordered_final_stream(hash_table, spill_context)?; timer.done(); - match result { - // Finished spilling the aggregate table, continue aggregating from input. - Ok(()) => ControlFlow::Continue(FinalHashAggregateState::ReadingInput { - hash_table, - spill_context: Some(spill_context), - }), - Err(e) => Self::break_with_err(e), + // Forwards output from the fully ordered stream that consumes the merged + // spill runs. + // + // Not wrapping in a timer and not record output batches since this is now `merge_stream` responsibility + // we just pass through + while let Some(batch) = output_stream.next().await.transpose()? { + emitter.emit(batch).await; } + + Ok(()) } - /// 1. Spills the last in-memory run. - /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// 1. Constructs a globally ordered input stream by applying a sort-preserving /// merge to all spills. - /// 3. Constructs a replay stream: an ordered final aggregate stream over the + /// 2. Constructs a replay stream: an ordered final aggregate stream over the /// fully ordered input constructed from the spills. /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_preparing_merge_input( + /// Returns the replay stream + fn switch_to_ordered_final_stream( &mut self, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::PreparingMergeInput { - mut hash_table, - mut spill_context, - } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected PreparingMergeInput state", - ); - }; - - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let replay = match spill_context.spill_table(&mut hash_table) { - Ok(()) => { - let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); - drop(hash_table); - match self.reservation.try_resize(0) { - Ok(()) => (*spill_context).into_replay_stream( - &self.baseline_metrics, - metrics, - self.reservation.new_empty(), - ), - Err(e) => Err(e), - } - } - Err(e) => Err(e), - }; - timer.done(); - - match replay { - Ok(stream) => { - ControlFlow::Continue(FinalHashAggregateState::MergingSpills { stream }) - } - Err(e) => Self::break_with_err(e), - } - } - - /// Forwards output from the fully ordered stream that consumes the merged - /// spill runs. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_merging_spills( - &mut self, - cx: &mut Context<'_>, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::MergingSpills { mut stream } = original_state else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected MergingSpills state", - ); - }; - - match stream.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - FinalHashAggregateState::MergingSpills { stream }, - )), - Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( - Poll::Ready(Some(Ok(batch))), - FinalHashAggregateState::MergingSpills { stream }, - )), - Poll::Ready(Some(Err(e))) => Self::break_with_err(e), - Poll::Ready(None) => ControlFlow::Continue(FinalHashAggregateState::Done), - } + hash_table: AggregateHashTable, + spill_context: Box, + ) -> Result { + let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); + drop(hash_table); + self.reservation.try_resize(0)?; + spill_context.into_replay_stream( + &self.baseline_metrics, + metrics, + self.reservation.new_empty(), + ) } - /// Handle ProducingOutput state - emit final aggregate value batches. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_producing_output( + /// Emit final aggregate value batches: + /// Input was exhausted without spilling, or the soft group limit was reached. + async fn produce_output_from_memory( &mut self, - original_state: FinalHashAggregateState, - ) -> FinalHashAggregateStateTransition { - let FinalHashAggregateState::ProducingOutput { mut hash_table } = original_state - else { - return Self::break_with_internal_err( - "Final hash aggregate stream expected ProducingOutput state", - ); - }; - + mut hash_table: AggregateHashTable, + mut emitter: TryEmitter, + ) -> Result<()> { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.next_output_batch(); - timer.done(); - - match result { - Ok(Some(batch)) => { - let next_state = if hash_table.is_done() { - drop(hash_table); - if let Err(e) = self.reservation.try_resize(0) { - return Self::break_with_err(e); - } - FinalHashAggregateState::Done - } else { - if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) - { - return Self::break_with_err(e); - } - FinalHashAggregateState::ProducingOutput { hash_table } - }; - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) - } - Err(e) => Self::break_with_err(e), - Ok(None) => { - drop(hash_table); - let next_state = FinalHashAggregateState::Done; - if let Err(e) = self.reservation.try_resize(0) { - return Self::break_with_err(e); - } - ControlFlow::Continue(next_state) - } - } - } -} - -impl Stream for FinalHashAggregateStream { - type Item = Result; + let mut timer = elapsed_compute.timer(); + hash_table.start_output()?; - /// Entry point for the final hash aggregate state machine. - /// - /// See comments in [`FinalHashAggregateStream`] for high-level ideas. - /// - /// State transition graph: - /// - /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling partial-state input and aggregating - /// those states into the final hash table. - /// - /// ReadingInput - /// -> ReadingInput - /// Aggregate one partial-state input batch. If it fits in memory, - /// continue with the next input batch. - /// -> Spilling - /// The table cannot reserve enough memory. Move all current states into - /// one fully group-key-sorted spill run. - /// -> ProducingOutput - /// Input was exhausted without spilling, or the soft group limit was - /// reached. Start outputting final aggregate values. - /// -> PreparingMergeInput - /// Input was exhausted after spilling. Spill the last in-memory run and - /// construct the ordered input used to merge all spill files. - /// - /// Spilling - /// -> ReadingInput - /// One sorted run was written; resume reading the original input. - /// - /// PreparingMergeInput - /// Spill the final in-memory run and build the input ordered replay stream. - /// -> MergingSpills - /// The final run was spilled and the ordered replay stream was built. - /// - /// MergingSpills - /// Aggregate the merged spill runs and emit final results. - /// -> MergingSpills - /// Forward one result batch from the fully ordered replay stream that - /// consumes the sort-preserving merge. - /// -> Done - /// The merged spill input was fully aggregated. - /// - /// ProducingOutput - /// -> ProducingOutput - /// One final output batch was yielded; repeat to continue producing - /// output incrementally. - /// -> Done - /// All final output was emitted. - /// - /// Any active state - /// -> Error - /// An error drops state-owned resources before it is returned. - /// - /// Error - /// -> (end) - /// - /// Done - /// -> (end) - /// ``` - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { loop { - let cur_state = self - .state - .take() - .expect("FinalHashAggregateStream state should not be None"); - - let next_state = match cur_state { - state @ FinalHashAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ FinalHashAggregateState::Spilling { .. } => { - self.handle_spilling(state) - } - state @ FinalHashAggregateState::PreparingMergeInput { .. } => { - self.handle_preparing_merge_input(state) - } - state @ FinalHashAggregateState::MergingSpills { .. } => { - self.handle_merging_spills(cx, state) - } - state @ FinalHashAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(state) - } - state @ FinalHashAggregateState::Error => { - self.close_input(); - self.reservation.free(); - self.state = Some(state); - return Poll::Ready(None); - } - state @ FinalHashAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Poll::Ready(None); - } + let Some(batch) = hash_table.next_output_batch()? else { + // Only reachable when the table held no groups at all: a + // non-empty table always reports its last batch together with + // the `Done` state, which the `try_resize` below already zeroes. + self.reservation.try_resize(0)?; + return Ok(()); }; - match next_state { - ControlFlow::Continue(next_state) => { - self.state = Some(next_state); - } - ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { - debug_assert!(matches!(next_state, FinalHashAggregateState::Error)); + // The table hands over its groups as they are materialized and + // reports a size of 0 once it reaches `Done`, so this releases the + // reservation before the final batch goes downstream. + self.reservation.try_resize(hash_table.memory_size())?; - // The handler has already discarded its state-owned resources. - // Release the remaining stream-owned resources before returning. - self.close_input(); - self.reservation.free(); - self.state = Some(FinalHashAggregateState::Error); - return Poll::Ready(Some(Err(e))); - } - ControlFlow::Break((poll, next_state)) => { - self.state = Some(next_state); - return poll; - } - } + timer.done(); + emitter + .emit(batch.record_output(&self.baseline_metrics)) + .await; + timer = elapsed_compute.timer(); } } } -impl RecordBatchStream for FinalHashAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index f00f1a160e27a..a6b5819916275 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -187,7 +187,6 @@ use datafusion_common::{ assert_eq_or_internal_err, internal_err, not_impl_err, }; use datafusion_execution::TaskContext; -use datafusion_execution::memory_pool::MemoryLimit; use datafusion_expr::{Accumulator, Aggregate}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -716,7 +715,7 @@ impl From for SendableRecordBatchStream { StreamType::AggregateStream(stream) => Box::pin(stream), StreamType::PartialHash(stream) => Box::pin(stream), StreamType::PartialReduceHash(stream) => Box::pin(stream), - StreamType::FinalHash(stream) => Box::pin(stream), + StreamType::FinalHash(stream) => stream.into_stream(), StreamType::SingleHash(stream) => Box::pin(stream), StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), @@ -1308,12 +1307,7 @@ impl AggregateExec { && self.group_by.is_single() } - fn should_use_partial_reduce_hash_stream(&self, context: &TaskContext) -> bool { - // TODO: implement memory-limited path and remove this limitation - if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { - return false; - } - + fn should_use_partial_reduce_hash_stream(&self, _context: &TaskContext) -> bool { self.mode == AggregateMode::PartialReduce && self.limit_options.is_none() && self.input_order_mode == InputOrderMode::Linear @@ -3241,7 +3235,7 @@ mod tests { use arrow::array::{ BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, - Int64Array, StructArray, UInt32Array, UInt64Array, + Int64Array, NullArray, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; @@ -4586,6 +4580,15 @@ mod tests { } fn partial_reduce_test_aggregate() -> Result { + partial_reduce_test_aggregate_with_batches(1) + } + + /// Partial-reduce aggregate over `num_input_batches` identical input batches + /// of partial states, each reducing to groups `1, 2, 3` with sums + /// `50, 20, 30`. + fn partial_reduce_test_aggregate_with_batches( + num_input_batches: usize, + ) -> Result { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Float64, false), @@ -4618,7 +4621,7 @@ mod tests { ], )?; let partial_reduce_input = TestMemoryExec::try_new_exec( - &[vec![partial_state_batch]], + &[vec![partial_state_batch; num_input_batches]], Arc::clone(&partial_schema), None, )?; @@ -4654,26 +4657,162 @@ mod tests { Ok(()) } - /// Spilling behavior is not implemented for partial-reduce stream yet, so fall - /// back to the existing `GroupedHashAggregateStream` + /// Partial-reduce hash aggregation emits its accumulated partial states early + /// under memory pressure instead of failing, and the early-emitted states + /// still merge into the correct result. #[tokio::test] - async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { - let partial_reduce = partial_reduce_test_aggregate()?; + async fn partial_reduce_aggregate_with_memory_limit_emits_early() -> Result<()> { + let num_input_batches = 3; + let partial_reduce = + partial_reduce_test_aggregate_with_batches(num_input_batches)?; let runtime = RuntimeEnvBuilder::new() .with_memory_limit(1, 1.0) .build_arc()?; - let task_ctx = - Arc::new( - TaskContext::default() - .with_session_config(SessionConfig::new().set_bool( - "datafusion.execution.enable_migration_aggregate", - true, - )) - .with_runtime(runtime), - ); + // A batch size smaller than the number of flushed groups also covers + // splitting one flush across several output batches. + let batch_size = 2; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)) + .with_runtime(runtime), + ); let stream = partial_reduce.execute_typed(0, &task_ctx)?; - assert!(matches!(stream, StreamType::GroupedHash(_))); + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + + // The table is flushed after every input batch, so each of the three + // groups is emitted once per input batch instead of being merged into a + // single row. Each flush is sliced into batches of 2 and 1 rows. + assert_eq!(output.len(), 2 * num_input_batches); + assert_snapshot!(batches_to_string(&output), @r" + +---+-------------+ + | a | SUM(b)[sum] | + +---+-------------+ + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + | 1 | 50.0 | + | 2 | 20.0 | + | 3 | 30.0 | + +---+-------------+ + "); + + Ok(()) + } + + /// Same shape as [`partial_reduce_test_aggregate_with_batches`], but with multiple + /// group keys. + fn partial_reduce_test_aggregate_rows_multi_group_keys( + num_input_batches: usize, + ) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("n", DataType::Null, true), + Field::new("b", DataType::Float64, false), + ])); + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema)?, "a".to_string()), + (col("n", &schema)?, "n".to_string()), + ]); + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), + Arc::new(NullArray::new(4)), + Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let partial_reduce_input = TestMemoryExec::try_new_exec( + &[vec![partial_state_batch; num_input_batches]], + Arc::clone(&partial_schema), + None, + )?; + + AggregateExec::try_new( + AggregateMode::PartialReduce, + group_by, + aggregates, + vec![None], + partial_reduce_input, + partial_schema, + ) + } + + #[tokio::test] + async fn partial_reduce_aggregate_with_memory_limit_emits_early_multi_group_keys() + -> Result<()> { + let num_input_batches = 3; + let partial_reduce = + partial_reduce_test_aggregate_rows_multi_group_keys(num_input_batches)?; + + // Pin the representation: this is exactly the condition + // `new_group_values` uses to pick `GroupValuesRows` over + // `GroupValuesColumn`. If a `Null` `GroupColumn` is ever added, this + // assertion fires and the test stops covering the row-encoded path. + let group_schema = partial_reduce + .group_by + .group_schema(&partial_reduce.schema())?; + assert!( + !group_values::multi_group_by::supported_schema(&group_schema), + "expected the Null group column to force the GroupValuesRows fallback" + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(1, 1.0) + .build_arc()?; + let batch_size = 2; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)) + .with_runtime(runtime), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + + // Same flush cadence as the column-backed test: one flush per input + // batch, each sliced into batches of 2 and 1 rows. + assert_eq!(output.len(), 2 * num_input_batches); + assert_snapshot!(batches_to_string(&output), @r" + +---+---+-------------+ + | a | n | SUM(b)[sum] | + +---+---+-------------+ + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + | 1 | | 50.0 | + | 2 | | 20.0 | + | 3 | | 30.0 | + +---+---+-------------+ + "); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index 9eedf868477a0..d8f1447bc9521 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -28,7 +28,7 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::stream::{Stream, StreamExt}; @@ -68,6 +68,15 @@ use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; /// This stage is useful for tree-reduce plans. It consumes the same schema as /// a final aggregate stage, but emits the same schema as a partial aggregate /// stage. +/// +/// # Memory Management +/// +/// If the memory reservation cannot grow after aggregating an input batch, all +/// accumulated partial states are emitted immediately, and the remaining input +/// is aggregated with an empty table. This repeats until the input ends. +/// +/// See [`crate::aggregates::AggregateMode::PartialReduce`] for why it's allowed +/// to emit the same group multiple times. pub(crate) struct PartialReduceHashAggregateStream { /// Output schema: group columns followed by partial aggregate state columns. schema: SchemaRef, @@ -75,6 +84,9 @@ pub(crate) struct PartialReduceHashAggregateStream { /// Input batches containing partial aggregate state rows. input: SendableRecordBatchStream, + /// Target output batch size from configuration. + batch_size: usize, + /// Execution metrics shared with the aggregate plan node. baseline_metrics: BaselineMetrics, @@ -93,10 +105,22 @@ enum PartialReduceHashAggregateState { ReadingInput { hash_table: AggregateHashTable, }, + /// A fully materialized partial-state batch being emitted incrementally + /// because the table ran out of memory while reading input. + EmittingOnMemoryPressure { + hash_table: AggregateHashTable, + // After each incremental emitting step, `remaining_groups` is updated + // with batch slicing. + remaining_groups: RecordBatch, + }, ProducingOutput { hash_table: AggregateHashTable, }, Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, } type PartialReduceHashAggregatePoll = Poll>>; @@ -111,28 +135,34 @@ type PartialReduceHashAggregateStateTransition = ControlFlow< impl PartialReduceHashAggregateState { fn hash_table(&self) -> &AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } fn hash_table_mut(&mut self) -> &mut AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } fn into_hash_table(self) -> AggregateHashTable { match self { - Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { - hash_table + Self::ReadingInput { hash_table } + | Self::EmittingOnMemoryPressure { hash_table, .. } + | Self::ProducingOutput { hash_table } => hash_table, + Self::Done | Self::Error => { + unreachable!("Done and Error states do not hold a hash table") } - Self::Done => unreachable!("Done state does not hold a hash table"), } } @@ -178,19 +208,25 @@ impl PartialReduceHashAggregateStream { Ok(Self { schema, input, + batch_size, baseline_metrics, reservation, state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), }) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - ) -> Result<()> { + fn close_input(&mut self) { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - hash_table.start_output() + } + + fn break_with_err( + error: DataFusionError, + ) -> PartialReduceHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialReduceHashAggregateState::Error, + )) } /// Handle ReadingInput state - aggregate partial state batches into the hash table. @@ -219,46 +255,126 @@ impl PartialReduceHashAggregateStream { timer.done(); if let Err(e) = result { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); + return Self::break_with_err(e); } - if let Err(e) = self - .reservation - .try_resize(original_state.hash_table().memory_size()) - { - return ControlFlow::Break(( - Poll::Ready(Some(Err(e))), - original_state, - )); - } - - ControlFlow::Continue(original_state) - } - Poll::Ready(Some(Err(e))) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + // Update the memory reservation. If OOM, do early emit. + self.resize_or_emit_early(original_state) } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), // Input ends, move to output state Poll::Ready(None) => { + self.close_input(); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = self.start_output(original_state.hash_table_mut()); + let result = original_state.hash_table_mut().start_output(); timer.done(); match result { Ok(()) => { ControlFlow::Continue(original_state.into_producing_output()) } - Err(e) => { - ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) - } + Err(e) => Self::break_with_err(e), } } } } + /// Update the memory reservation. If the reservation succeeds, continue reading + /// input. If OOM, clear the aggregated states in the hash table, and early emit + /// them immediately. + /// + /// Returns the next state; the caller finishes the intended task based on it. + /// + /// The reservation is left at its pre-emission size while the states are being + /// emitted, because the cleared states are still held in memory as + /// `remaining_groups`. The reservation will be reset after exiting the + /// `EmittingOnMemoryPressure` state. + /// + /// # Implementation Note + /// All accumulated states are materialized at once, and then sliced into + /// `batch_size` output batches. Emit them incrementally after blocked state + /// management is ready. + /// + /// Issue: + fn resize_or_emit_early( + &mut self, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); // Stop on drop + let resize_result = self + .reservation + .try_resize(original_state.hash_table().memory_size()); + + let oom = match resize_result { + Ok(()) => return ControlFlow::Continue(original_state), + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Self::break_with_err(e), + }; + + let state_batch_result = original_state.hash_table_mut().take_state_batch(); + + match state_batch_result { + Ok(Some(remaining_groups)) => ControlFlow::Continue( + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table: original_state.into_hash_table(), + remaining_groups, + }, + ), + // No accumulated group to emit, so early emission cannot release any + // memory: report the original error. + Ok(None) => Self::break_with_err(oom), + Err(e) => Self::break_with_err(e), + } + } + + /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state + /// batch in `batch_size`(from configuration) slices. After all slices are + /// emitted, update the memory reservation and resume reading input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_emitting_on_memory_pressure( + &mut self, + original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + let PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: batch, + } = original_state + else { + unreachable!("expected the EmittingOnMemoryPressure state") + }; + + let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + // Go back to `ReadingInput` + ( + batch, + PartialReduceHashAggregateState::ReadingInput { hash_table }, + ) + } else { + // More batches to output, continue in the current state. + let remaining = + batch.slice(self.batch_size, batch.num_rows() - self.batch_size); + let output = batch.slice(0, self.batch_size); + ( + output, + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: remaining, + }, + ) + }; + + debug_assert!(output_batch.num_rows() > 0); + ControlFlow::Break(( + Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + /// Handle ProducingOutput state - emit merged partial aggregate state batches. /// /// See comments at `poll_next()` for details. @@ -281,6 +397,8 @@ impl PartialReduceHashAggregateStream { match result { Ok(Some(batch)) => { + // The output is already materialized, so a failed resize cannot + // be acted on: keep the reservation as is and finish the output. let _ = self .reservation .try_resize(original_state.hash_table().memory_size()); @@ -300,7 +418,7 @@ impl PartialReduceHashAggregateStream { let _ = self.reservation.try_resize(0); ControlFlow::Continue(original_state.into_done()) } - Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + Err(e) => Self::break_with_err(e), } } } @@ -325,10 +443,22 @@ impl Stream for PartialReduceHashAggregateStream { /// Aggregate one partial-state input batch, update the inner aggregate /// hash table, and continue with the next input batch. /// + /// -> EmittingOnMemoryPressure + /// The table cannot reserve enough memory. Materialize all accumulated + /// partial states and begin emitting them incrementally. + /// /// -> ProducingOutput /// Input was exhausted. Move to the next state to start outputting /// merged partial aggregate states. /// + /// EmittingOnMemoryPressure + /// -> EmittingOnMemoryPressure + /// One batch-sized slice was yielded; repeat until all materialized + /// partial states are emitted. + /// + /// -> ReadingInput + /// The materialized states were emitted; continue with the empty table. + /// /// ProducingOutput /// -> ProducingOutput /// One merged partial-state output batch was yielded; repeat to @@ -337,6 +467,13 @@ impl Stream for PartialReduceHashAggregateStream { /// -> Done /// All merged partial-state output was emitted. /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// /// Done /// -> (end) /// ``` @@ -354,9 +491,18 @@ impl Stream for PartialReduceHashAggregateStream { state @ PartialReduceHashAggregateState::ReadingInput { .. } => { self.handle_reading_input(cx, state) } + state @ PartialReduceHashAggregateState::EmittingOnMemoryPressure { + .. + } => self.handle_emitting_on_memory_pressure(state), state @ PartialReduceHashAggregateState::ProducingOutput { .. } => { self.handle_producing_output(state) } + state @ PartialReduceHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } state @ PartialReduceHashAggregateState::Done => { let _ = self.reservation.try_resize(0); self.state = Some(state); @@ -368,6 +514,19 @@ impl Stream for PartialReduceHashAggregateStream { ControlFlow::Continue(next_state) => { self.state = Some(next_state); } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!( + next_state, + PartialReduceHashAggregateState::Error + )); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(PartialReduceHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } ControlFlow::Break((poll, next_state)) => { self.state = Some(next_state); return poll; diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index f63418f63e5f0..8cd603e91eea5 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -1078,6 +1078,7 @@ async fn collect_left_input( while let Some(batch) = stream.next().await { let batch = batch?; + let build_timer = metrics.build_time.timer(); let batch_size = batch.get_array_memory_size(); match reservation.try_grow(batch_size) { Ok(()) => { @@ -1087,6 +1088,9 @@ async fn collect_left_input( batches.push(batch); } Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { + // Do not keep the operator timer running while the spill path + // drains the child stream. + build_timer.done(); let spill_manager = spill_manager.expect("checked by is_spillable_oom"); let spilled = spill_left_input( spill_manager, @@ -1109,6 +1113,10 @@ async fn collect_left_input( } } + // Only time the build-side materialization performed by this operator, not + // polling the child stream above. + let build_timer = metrics.build_time.timer(); + let merged_batch = concat_batches(&schema, &batches)?; // Reserve memory for visited_left_side bitmap if required by join type @@ -1118,6 +1126,9 @@ async fn collect_left_input( match reservation.try_grow(buffer_size) { Ok(()) => {} Err(e) if is_spillable_oom(&e, spill_manager.as_ref()) => { + // `spill_left_input` owns its timing and polls the input stream + // outside that timer. + build_timer.done(); let spill_manager = spill_manager.expect("checked by is_spillable_oom"); drop(batches); let spilled = spill_left_input( @@ -1198,6 +1209,7 @@ async fn spill_left_input( metrics: BuildProbeJoinMetrics, reservation: &MemoryReservation, ) -> Result> { + let build_timer = metrics.build_time.timer(); let mut spill_file = spill_manager.create_in_progress_file("NestedLoopJoin left spill")?; @@ -1217,9 +1229,11 @@ async fn spill_left_input( spill_file.append_batch(&batch)?; } } + build_timer.done(); while let Some(batch) = stream.next().await { let batch = batch?; + let _build_timer = metrics.build_time.timer(); if batch.num_rows() > 0 { metrics.build_input_batches.add(1); metrics.build_input_rows.add(batch.num_rows()); @@ -1227,6 +1241,7 @@ async fn spill_left_input( } } + let _build_timer = metrics.build_time.timer(); Ok(spill_file.finish()?.map(|file| LeftSpillData { spill_manager, spill_file: file, @@ -1524,12 +1539,6 @@ impl Stream for NestedLoopJoinStream { // side batch, before start joining. NLJState::BufferingLeft => { debug!("[NLJState] Entering: {:?}", self.state); - // inside `collect_left_input` (the routine to buffer build - // -side batches), related metrics except build time will be - // updated. - // stop on drop - let build_metric = self.metrics.join_metrics.build_time.clone(); - let _build_timer = build_metric.timer(); match self.handle_buffering_left(cx) { ControlFlow::Continue(()) => {} @@ -1562,9 +1571,6 @@ impl Stream for NestedLoopJoinStream { // handling (e.g., in cases like left join). NLJState::FetchingRight => { debug!("[NLJState] Entering: {:?}", self.state); - // stop on drop - let join_metric = self.metrics.join_metrics.join_time.clone(); - let _join_timer = join_metric.timer(); match self.handle_fetching_right(cx) { ControlFlow::Continue(()) => {} @@ -1682,9 +1688,6 @@ impl Stream for NestedLoopJoinStream { NLJState::EmitGlobalRightUnmatched => { debug!("[NLJState] Entering: {:?}", self.state); - let join_metric = self.metrics.join_metrics.join_time.clone(); - let _join_timer = join_metric.timer(); - match self.handle_emit_global_right_unmatched(cx) { ControlFlow::Continue(()) => {} ControlFlow::Break(poll) => { @@ -1879,6 +1882,7 @@ impl NestedLoopJoinStream { &mut self, cx: &mut std::task::Context<'_>, ) -> ControlFlow>>> { + let build_metric = self.metrics.join_metrics.build_time.clone(); let SpillState::Active(active) = &mut self.spill_state else { unreachable!( "handle_buffering_left_memory_limited called without Active spill state" @@ -1889,6 +1893,7 @@ impl NestedLoopJoinStream { // stream was consumed, open a fresh stream over the left spill file. if active.left_stream.is_none() { let spill_data = Arc::clone(&active.left_spill); + let _build_timer = build_metric.timer(); match spill_data .spill_manager .read_spill_as_stream(Arc::clone(&spill_data.spill_file), None) @@ -1914,6 +1919,7 @@ impl NestedLoopJoinStream { loop { match left_stream.poll_next_unpin(cx) { Poll::Ready(Some(Ok(batch))) => { + let _build_timer = build_metric.timer(); if batch.num_rows() == 0 { continue; } @@ -1954,6 +1960,8 @@ impl NestedLoopJoinStream { } } + let _build_timer = build_metric.timer(); + // If the left stream is fully exhausted, release its resources so the // upstream pipeline can be torn down before we move on to probing. if self.left_exhausted { @@ -2046,47 +2054,52 @@ impl NestedLoopJoinStream { &mut self, cx: &mut std::task::Context<'_>, ) -> ControlFlow>>> { - match self + let result = match self .right_data .as_mut() .expect("right_data must be present while fetching right") .poll_next_unpin(cx) { - Poll::Ready(result) => match result { - Some(Ok(right_batch)) => { - // Update metrics - let right_batch_rows = right_batch.num_rows(); - self.metrics.join_metrics.input_rows.add(right_batch_rows); - self.metrics.join_metrics.input_batches.add(1); - - // Skip the empty batch - if right_batch_rows == 0 { - return ControlFlow::Continue(()); - } + Poll::Ready(result) => result, + Poll::Pending => return ControlFlow::Break(Poll::Pending), + }; - self.current_right_batch = Some(right_batch); + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); - // Prepare right bitmap - if self.should_track_unmatched_right { - let zeroed_buf = BooleanBuffer::new_unset(right_batch_rows); - self.current_right_batch_matched = - Some(BooleanArray::new(zeroed_buf, None)); - } + match result { + Some(Ok(right_batch)) => { + // Update metrics + let right_batch_rows = right_batch.num_rows(); + self.metrics.join_metrics.input_rows.add(right_batch_rows); + self.metrics.join_metrics.input_batches.add(1); - self.left_probe_idx = 0; - self.state = NLJState::ProbeRight; - ControlFlow::Continue(()) + // Skip the empty batch + if right_batch_rows == 0 { + return ControlFlow::Continue(()); } - Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), - None => { - // Right side exhausted: probing for the current left chunk - // is finished. `ProbeEnd` reports probe completion before - // emitting unmatched-left rows. - self.state = NLJState::ProbeEnd; - ControlFlow::Continue(()) + + self.current_right_batch = Some(right_batch); + + // Prepare right bitmap + if self.should_track_unmatched_right { + let zeroed_buf = BooleanBuffer::new_unset(right_batch_rows); + self.current_right_batch_matched = + Some(BooleanArray::new(zeroed_buf, None)); } - }, - Poll::Pending => ControlFlow::Break(Poll::Pending), + + self.left_probe_idx = 0; + self.state = NLJState::ProbeRight; + ControlFlow::Continue(()) + } + Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + None => { + // Right side exhausted: probing for the current left chunk + // is finished. `ProbeEnd` reports probe completion before + // emitting unmatched-left rows. + self.state = NLJState::ProbeEnd; + ControlFlow::Continue(()) + } } } @@ -2303,6 +2316,8 @@ impl NestedLoopJoinStream { // On first entry, open a new replay pass on the right input if self.right_data.is_none() { + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); let SpillState::Active(ref mut active) = self.spill_state else { unreachable!("EmitGlobalRightUnmatched without Active spill state"); }; @@ -2318,13 +2333,20 @@ impl NestedLoopJoinStream { } // Poll the replay stream for the next right batch - match self + let result = match self .right_data .as_mut() .expect("right_data must be present") .poll_next_unpin(cx) { - Poll::Ready(Some(Ok(right_batch))) => { + Poll::Ready(result) => result, + Poll::Pending => return ControlFlow::Break(Poll::Pending), + }; + + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); + match result { + Some(Ok(right_batch)) => { if right_batch.num_rows() == 0 { return ControlFlow::Continue(()); } @@ -2370,8 +2392,8 @@ impl NestedLoopJoinStream { Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), } } - Poll::Ready(Some(Err(e))) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), - Poll::Ready(None) => { + Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + None => { // All right batches replayed match self.output_buffer.finish_buffered_batch() { Ok(()) => { @@ -2381,7 +2403,6 @@ impl NestedLoopJoinStream { Err(e) => ControlFlow::Break(Poll::Ready(Some(arrow_err!(e)))), } } - Poll::Pending => ControlFlow::Break(Poll::Pending), } } @@ -3378,6 +3399,9 @@ fn build_unmatched_batch( #[cfg(test)] pub(crate) mod tests { + use std::pin::Pin; + use std::time::Duration; + use super::*; use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::{TestMemoryExec, assert_join_metrics}; @@ -3387,9 +3411,15 @@ pub(crate) mod tests { use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field}; + use bytes::Bytes; use datafusion_common::assert_contains; + use datafusion_common::instant::Instant; use datafusion_common::test_util::batches_to_sort_string; + use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, + }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; use datafusion_physical_expr::{Partitioning, PhysicalExpr}; @@ -3399,6 +3429,73 @@ pub(crate) mod tests { use insta::assert_snapshot; use rstest::rstest; + fn delayed_stream(batch: RecordBatch, delay: Duration) -> SendableRecordBatchStream { + let schema = batch.schema(); + Box::pin(crate::stream::RecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { + std::thread::sleep(delay); + Ok(batch) + }), + )) + } + + /// Delays the first item while the spill stream is polled, making an + /// incorrectly scoped operator timer include the delay. + struct DelayedReadSpillFile { + inner: Arc, + delay: Duration, + read_count: Arc, + } + + impl SpillFile for DelayedReadSpillFile { + fn path(&self) -> Option<&std::path::Path> { + self.inner.path() + } + + fn size(&self) -> Option { + self.inner.size() + } + + fn read_stream( + &self, + ) -> Result> + Send>>> { + let delay = self.delay; + let read_count = Arc::clone(&self.read_count); + let mut delay_first_item = true; + let stream = self.inner.read_stream()?.map(move |item| { + if delay_first_item { + delay_first_item = false; + read_count.fetch_add(1, Ordering::Relaxed); + std::thread::sleep(delay); + } + item + }); + Ok(Box::pin(stream)) + } + + fn open_writer(&self) -> Result> { + self.inner.open_writer() + } + } + + /// Wraps local spill files so replay reads can be delayed deterministically. + struct DelayedReadTempFileFactory { + inner: Arc, + delay: Duration, + read_count: Arc, + } + + impl TempFileFactory for DelayedReadTempFileFactory { + fn create_temp_file(&self, description: &str) -> Result> { + Ok(Arc::new(DelayedReadSpillFile { + inner: self.inner.create_tmp_file(description)?, + delay: self.delay, + read_count: Arc::clone(&self.read_count), + })) + } + } + fn build_table( a: (&str, &Vec), b: (&str, &Vec), @@ -3554,6 +3651,246 @@ pub(crate) mod tests { ) } + async fn run_join_with_child_poll_delays( + left_delay: Duration, + right_delay: Duration, + memory_limited: bool, + ) -> Result<(Duration, Duration, Duration)> { + run_join_with_poll_delays( + left_delay, + right_delay, + memory_limited, + JoinType::Inner, + None, + ) + .await + } + + async fn run_join_with_poll_delays( + left_delay: Duration, + right_delay: Duration, + memory_limited: bool, + join_type: JoinType, + spill_read_delay: Option<(Duration, Arc)>, + ) -> Result<(Duration, Duration, Duration)> { + let left_batch = + build_table_i32(("a1", &vec![1]), ("b1", &vec![2]), ("c1", &vec![3])); + let right_batch = + build_table_i32(("a2", &vec![4]), ("b2", &vec![5]), ("c2", &vec![6])); + let left_schema = left_batch.schema(); + let right_schema = right_batch.schema(); + + let left_stream = delayed_stream(left_batch.clone(), left_delay); + let right_stream = delayed_stream(right_batch, right_delay); + + let task_ctx = if let Some((delay, read_count)) = spill_read_delay { + let inner = Arc::new( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .build()?, + ); + let runtime = RuntimeEnvBuilder::new() + .with_disk_manager_builder(DiskManagerBuilder::default().with_mode( + DiskManagerMode::Custom(Arc::new(DelayedReadTempFileFactory { + inner, + delay, + read_count, + })), + )) + .build_arc()?; + Arc::new(TaskContext::default().with_runtime(runtime)) + } else { + Arc::new(TaskContext::default()) + }; + let metrics_set = ExecutionPlanMetricsSet::new(); + let metrics = NestedLoopJoinMetrics::new(&metrics_set, 0); + let build_time = metrics.join_metrics.build_time.clone(); + let join_time = metrics.join_metrics.join_time.clone(); + let (left_data, right_stream, spill_state) = if memory_limited { + let reservation = MemoryConsumer::new("NestedLoopJoinLoad[test]".to_string()) + .with_can_spill(true) + .register(task_ctx.memory_pool()); + let global_right_bitmaps_reservation = + MemoryConsumer::new("NestedLoopJoinGlobalRightBitmaps[test]".to_string()) + .register(task_ctx.memory_pool()); + let spill_manager = SpillManager::new( + task_ctx.runtime_env(), + metrics.spill_metrics.clone(), + Arc::clone(&right_schema), + ); + let left_spill_manager = SpillManager::new( + task_ctx.runtime_env(), + metrics.spill_metrics.clone(), + Arc::clone(&left_schema), + ); + let mut left_spill_file = + left_spill_manager.create_in_progress_file("test left spill")?; + left_spill_file.append_batch(&left_batch)?; + let left_spill_file = left_spill_file + .finish()? + .expect("the test left spill contains one batch"); + let active = SpillStateActive { + left_spill: Arc::new(LeftSpillData { + spill_manager: left_spill_manager, + spill_file: left_spill_file, + schema: Arc::clone(&left_schema), + }), + left_stream: Some(left_stream), + left_schema: Some(Arc::clone(&left_schema)), + reservation, + pending_batches: Vec::new(), + right_input: ReplayableStreamSource::new( + right_stream, + spill_manager, + "test right spill", + ), + global_right_bitmaps: Vec::new(), + global_right_bitmaps_reservation, + right_batch_index: 0, + }; + ( + OnceFut::new(async { internal_err!("unused left data was polled") }), + Box::pin(crate::EmptyRecordBatchStream::new(Arc::clone( + &right_schema, + ))) as SendableRecordBatchStream, + SpillState::Active(Box::new(active)), + ) + } else { + let reservation = MemoryConsumer::new("NestedLoopJoinLoad[test]".to_string()) + .register(task_ctx.memory_pool()); + ( + OnceFut::new(collect_left_input( + left_stream, + metrics.join_metrics.clone(), + reservation, + false, + 1, + None, + )), + right_stream, + SpillState::Disabled, + ) + }; + let (output_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + let stream = NestedLoopJoinStream::new( + Arc::new(output_schema), + None, + join_type, + right_stream, + left_data, + column_indices, + metrics, + 1024, + spill_state, + ); + + let start = Instant::now(); + let batches = common::collect(Box::pin(stream)).await?; + let wall_time = start.elapsed(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + + Ok(( + Duration::from_nanos(build_time.value() as u64), + Duration::from_nanos(join_time.value() as u64), + wall_time, + )) + } + + async fn check_child_poll_time_excluded(mut run: F) -> Result<()> + where + F: FnMut(Duration) -> Fut, + Fut: Future>, + { + // Escalating the delay filters out fixed-size scheduler preemption without + // masking the bug: an incorrectly scoped timer grows with every delay. + let mut delay = Duration::from_millis(50); + for attempt in 0..3 { + let (operator_time, wall_time) = run(delay).await?; + assert!( + !operator_time.is_zero(), + "operator work should still be timed" + ); + assert!( + wall_time >= delay, + "child poll delay should dominate wall time: {wall_time:?} < {delay:?}" + ); + if operator_time < delay { + return Ok(()); + } + assert!( + attempt < 2, + "operator time ({operator_time:?}) included the child poll delay ({delay:?})" + ); + delay *= 4; + } + unreachable!() + } + + #[tokio::test] + async fn build_time_excludes_left_child_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (build_time, _, wall_time) = + run_join_with_child_poll_delays(delay, Duration::ZERO, false).await?; + Ok((build_time, wall_time)) + }) + .await + } + + #[tokio::test] + async fn join_time_excludes_right_child_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (_, join_time, wall_time) = + run_join_with_child_poll_delays(Duration::ZERO, delay, false).await?; + Ok((join_time, wall_time)) + }) + .await + } + + #[tokio::test] + async fn build_time_excludes_spill_stream_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (build_time, _, wall_time) = + run_join_with_child_poll_delays(delay, Duration::ZERO, true).await?; + Ok((build_time, wall_time)) + }) + .await + } + + #[tokio::test] + async fn join_time_excludes_replayable_input_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let (_, join_time, wall_time) = + run_join_with_child_poll_delays(Duration::ZERO, delay, true).await?; + Ok((join_time, wall_time)) + }) + .await + } + + #[tokio::test(flavor = "current_thread")] + async fn join_time_excludes_global_right_unmatched_replay_poll() -> Result<()> { + check_child_poll_time_excluded(|delay| async move { + let read_count = Arc::new(AtomicUsize::new(0)); + // The only left chunk is supplied directly, so the sole spill read + // is the final right replay in EmitGlobalRightUnmatched. + let (_, join_time, wall_time) = run_join_with_poll_delays( + Duration::ZERO, + Duration::ZERO, + true, + JoinType::Right, + Some((delay, Arc::clone(&read_count))), + ) + .await?; + assert_eq!( + read_count.load(Ordering::Relaxed), + 1, + "EmitGlobalRightUnmatched should replay the right spill exactly once" + ); + Ok((join_time, wall_time)) + }) + .await + } + fn prepare_join_filter() -> JoinFilter { let column_indices = vec![ ColumnIndex { diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index a1a1ff6f04fe4..1eac6f974beed 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -429,7 +429,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::Cast(Cast::new_from_field(expr, Arc::new(field)))) } ExprType::TryCast(cast) => { @@ -442,7 +443,8 @@ pub fn parse_expr( let data_type: DataType = cast.arrow_type.as_ref().required("arrow_type")?; let field = data_type .into_nullable_field() - .with_nullable(cast.nullable.unwrap_or(true)); + .with_nullable(cast.nullable.unwrap_or(true)) + .with_metadata(cast.metadata.clone()); Ok(Expr::TryCast(TryCast::new_from_field( expr, Arc::new(field), diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 7c916db9c3cd1..750b20323ad2e 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -2693,6 +2693,19 @@ fn roundtrip_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::Cast(Cast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); } #[test] @@ -2703,6 +2716,19 @@ fn roundtrip_try_cast() { let ctx = SessionContext::new(); roundtrip_expr_test(test_expr, ctx); + let field = + Field::new("", DataType::Boolean, false).with_metadata(HashMap::from([( + String::from("key"), + String::from("value"), + )])); + let test_expr = Expr::TryCast(TryCast::new_from_field( + Box::new(lit(1.0_f32)), + Arc::new(field), + )); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); + let test_expr = Expr::TryCast(TryCast::new(Box::new(lit("not a bool")), DataType::Boolean)); diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs index 7418d0b5b7605..c335d3ee4c16d 100644 --- a/datafusion/sql/src/unparser/ast.rs +++ b/datafusion/sql/src/unparser/ast.rs @@ -167,6 +167,8 @@ pub struct SelectBuilder { /// Table aliases that correspond to LATERAL FLATTEN relations. /// Column references into these aliases must use `VALUE` as the column name. flatten_table_aliases: Vec, + /// Depth of explicitly named subqueries currently being rendered. + subquery_alias_depth: usize, } /// Prefix used for auto-generated LATERAL FLATTEN table aliases. @@ -195,6 +197,19 @@ impl SelectBuilder { self.flatten_table_aliases.iter().any(|a| a == alias) } + pub(super) fn enter_subquery_alias(&mut self) { + self.subquery_alias_depth += 1; + } + + pub(super) fn exit_subquery_alias(&mut self) { + debug_assert!(self.subquery_alias_depth > 0); + self.subquery_alias_depth -= 1; + } + + pub(super) fn inside_subquery_alias(&self) -> bool { + self.subquery_alias_depth > 0 + } + /// Returns the most recently generated flatten alias, or `None` if /// `next_flatten_alias` has not been called yet. pub fn current_flatten_alias(&self) -> Option { @@ -419,6 +434,7 @@ impl SelectBuilder { flavor: Some(SelectFlavor::Standard), flatten_alias_counter: 0, flatten_table_aliases: Vec::new(), + subquery_alias_depth: 0, } } } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 30320acbb24dc..1a4a094e96ed6 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -54,7 +54,7 @@ use datafusion_expr::{ TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias, }; use sqlparser::ast::{self, Ident, OrderByKind, SetExpr, TableAliasColumnDef}; -use std::{sync::Arc, vec}; +use std::{collections::HashSet, sync::Arc, vec}; /// Convert a DataFusion [`LogicalPlan`] to [`ast::Statement`] /// @@ -164,6 +164,12 @@ impl<'a> UnparserAggScope<'a> { } } +#[derive(Clone, Copy)] +struct DerivedInputScope<'a> { + alias: &'static str, + schema: &'a DFSchema, +} + impl Unparser<'_> { pub fn plan_to_sql(&self, plan: &LogicalPlan) -> Result { let mut plan = normalize_union_schema(plan)?; @@ -451,6 +457,102 @@ impl Unparser<'_> { } } + /// Return the alias recursion would assign when `plan` must become a + /// derived relation below an already rendered projection. + fn derived_input_alias(plan: &LogicalPlan) -> Option<&'static str> { + match plan { + LogicalPlan::Projection(_) => Some("derived_projection"), + LogicalPlan::Limit(_) => Some("derived_limit"), + LogicalPlan::Sort(_) => Some("derived_sort"), + LogicalPlan::Distinct(_) => Some("derived_distinct"), + LogicalPlan::Filter(filter) => { + Self::derived_input_alias(filter.input.as_ref()) + } + LogicalPlan::Repartition(repartition) => { + Self::derived_input_alias(repartition.input.as_ref()) + } + _ => None, + } + } + + fn derived_input_scope<'a>( + plan: &'a LogicalPlan, + select: &SelectBuilder, + ) -> Option> { + if select.inside_subquery_alias() { + return None; + } + + if select.already_projected() + && find_unnest_node_within_select(plan).is_none() + && let Some(alias) = Self::derived_input_alias(plan) + { + return Some(DerivedInputScope { + alias, + schema: plan.schema().as_ref(), + }); + } + + match plan { + LogicalPlan::Projection(projection) => { + let alias = Self::derived_input_alias(projection.input.as_ref())?; + let qualified_projection = projection.expr.iter().any(|expr| { + expr.column_refs() + .iter() + .any(|column| column.relation.is_some()) + }); + let mut input_names = HashSet::new(); + let unique_input_names = projection + .input + .schema() + .fields() + .iter() + .all(|field| input_names.insert(field.name())); + + (qualified_projection + && unique_input_names + && find_unnest_node_within_select(plan).is_none()) + .then_some(DerivedInputScope { + alias, + schema: projection.input.schema().as_ref(), + }) + } + LogicalPlan::Filter(filter) => { + Self::derived_input_scope(filter.input.as_ref(), select) + } + LogicalPlan::Limit(limit) => { + Self::derived_input_scope(limit.input.as_ref(), select) + } + LogicalPlan::Sort(sort) => { + Self::derived_input_scope(sort.input.as_ref(), select) + } + LogicalPlan::Repartition(repartition) => { + Self::derived_input_scope(repartition.input.as_ref(), select) + } + _ => None, + } + } + + fn rebase_derived_input_expr( + &self, + expr: Expr, + scope: Option>, + ) -> Result { + let Some(scope) = scope else { + return Ok(expr); + }; + if self.dialect.requires_derived_table_alias() { + let mut alias_rewriter = TableAliasRewriter { + table_schema: scope.schema, + alias_name: TableReference::bare(scope.alias), + rewrite_unqualified: false, + }; + expr.rewrite(&mut alias_rewriter).data() + } else { + Self::strip_column_qualifiers_for_schema(expr, scope.schema) + } + } + fn contains_aggregate_before_relation(plan: &LogicalPlan) -> bool { match plan { LogicalPlan::Aggregate(_) => true, @@ -829,6 +931,26 @@ impl Unparser<'_> { columns, ); } + + if let Some(scope) = Self::derived_input_scope(plan, select) { + // The input is about to enter a new SQL scope. Preserve that + // boundary explicitly and make the outer expressions resolve + // against the relation that will actually be visible there. + let requires_alias = self.dialect.requires_derived_table_alias(); + let alias = requires_alias + .then(|| self.new_table_alias(scope.alias.to_string(), vec![])); + self.derive(p.input.as_ref(), relation, alias, false)?; + + let items = p + .expr + .iter() + .cloned() + .map(|expr| self.rebase_derived_input_expr(expr, Some(scope))) + .map(|expr| self.select_item_to_sql(&expr?)) + .collect::>>()?; + select.projection(items); + return Ok(()); + } // For Snowflake FLATTEN: when the outer Projection has // UNNEST(...) display-name columns (from SELECT * / SELECT // UNNEST(...)), generate a flatten alias now so that @@ -1072,6 +1194,7 @@ impl Unparser<'_> { self.select_to_sql_recursively(cur, query, select, relation) } LogicalPlan::Filter(filter) => { + let derived_input_scope = Self::derived_input_scope(plan, select); let window = find_window_nodes_within_select( plan, None, @@ -1088,15 +1211,23 @@ impl Unparser<'_> { unprojected = UnparserAggScope::new(agg).prepare(unprojected, None)?; } + unprojected = + self.rebase_derived_input_expr(unprojected, derived_input_scope)?; let filter_expr = self.expr_to_sql(&unprojected)?; select.qualify(Some(filter_expr)); } else if let Some(agg) = agg { - let unprojected = UnparserAggScope::new(agg) + let mut unprojected = UnparserAggScope::new(agg) .prepare(filter.predicate.clone(), None)?; + unprojected = + self.rebase_derived_input_expr(unprojected, derived_input_scope)?; let filter_expr = self.expr_to_sql(&unprojected)?; select.having(Some(filter_expr)); } else { - let filter_expr = self.expr_to_sql(&filter.predicate)?; + let predicate = self.rebase_derived_input_expr( + filter.predicate.clone(), + derived_input_scope, + )?; + let filter_expr = self.expr_to_sql(&predicate)?; select.selection(Some(filter_expr)); } @@ -1172,17 +1303,25 @@ impl Unparser<'_> { )))); } + let derived_input_scope = Self::derived_input_scope(plan, select); let agg = find_agg_node_within_select(plan, select.already_projected()); // unproject sort expressions let sort_exprs: Vec = sort .expr .iter() .map(|sort_expr| { - Self::unproject_sort_expr_in_scope( + let sort_expr = Self::unproject_sort_expr_in_scope( sort_expr.clone(), agg, sort.input.as_ref(), - ) + )?; + Ok(SortExpr { + expr: self.rebase_derived_input_expr( + sort_expr.expr, + derived_input_scope, + )?, + ..sort_expr + }) }) .collect::>>()?; @@ -1537,7 +1676,8 @@ impl Unparser<'_> { )]); } let plan = unparsed_table_scan.unwrap_or_else(|| plan.clone()); - if !columns.is_empty() + select.enter_subquery_alias(); + let recursive_result = if !columns.is_empty() && !self.dialect.supports_column_alias_in_table_alias() { // Instead of specifying column aliases as part of the outer table, inject them directly into the inner projection @@ -1558,10 +1698,12 @@ impl Unparser<'_> { query, select, relation, - )?; + ) } else { - self.select_to_sql_recursively(&plan, query, select, relation)?; - } + self.select_to_sql_recursively(&plan, query, select, relation) + }; + select.exit_subquery_alias(); + recursive_result?; relation.alias(Some( self.new_table_alias(plan_alias.alias.table().to_string(), columns), diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index d6c31570bf1b0..32dfade056dd2 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -392,13 +392,94 @@ fn roundtrip_statement_with_dialect_4() -> Result<(), DataFusionError> { Ok(()) } +#[test] +fn roundtrip_rebases_derived_projection_references() -> Result<(), DataFusionError> { + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta);", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta)", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta);", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection`", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) order by j1_id;", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta) ORDER BY j1_id ASC NULLS LAST", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) order by j1_id;", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` ORDER BY `derived_projection`.`j1_id` ASC", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) where j1_id > 1;", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT ta.j1_id FROM j1 AS ta WHERE (ta.j1_id > 1))", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select ta.j1_id as j1_id from j1 ta) where j1_id > 1;", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta` WHERE (`ta`.`j1_id` > 1)) AS `derived_projection`", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select distinct ta.j1_id as j1_id from j1 ta);", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT j1_id FROM (SELECT DISTINCT ta.j1_id FROM j1 AS ta)", + ); + roundtrip_statement_with_dialect_helper!( + sql: "select j1_id from (select distinct ta.j1_id as j1_id from j1 ta);", + parser_dialect: MySqlDialect {}, + unparser_dialect: UnparserMySqlDialect {}, + expected: @"SELECT `derived_distinct`.`j1_id` FROM (SELECT DISTINCT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_distinct`", + ); + + let statement = Parser::new(&GenericDialect {}) + .try_with_sql("select j1_id from (select ta.j1_id as j1_id from j1 ta)")? + .parse_statement()?; + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let plan = SqlToRel::new(&context).sql_statement_to_plan(statement)?; + let plan = LogicalPlanBuilder::from(plan) + .filter(col("ta.j1_id").gt(lit(1)))? + .build()?; + let unparser = Unparser::new(&UnparserMySqlDialect {}); + assert_snapshot!( + unparser.plan_to_sql(&plan)?, + @"SELECT `derived_projection`.`j1_id` FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` WHERE (`derived_projection`.`j1_id` > 1)" + ); + + let schema = Schema::new(vec![Field::new("j1_id", DataType::Int32, false)]); + let plan = table_scan(Some("j1"), &schema, None)? + .alias("ta")? + .project(vec![col("ta.j1_id")])? + .filter(col("ta.j1_id").gt(lit(0)))? + .project(vec![lit(1)])? + .build()?; + assert_snapshot!( + unparser.plan_to_sql(&plan)?, + @"SELECT 1 FROM (SELECT `ta`.`j1_id` FROM `j1` AS `ta`) AS `derived_projection` WHERE (`derived_projection`.`j1_id` > 0)" + ); + Ok(()) +} + #[test] fn roundtrip_statement_with_dialect_5() -> Result<(), DataFusionError> { roundtrip_statement_with_dialect_helper!( sql: "select j1_id from (select j1_id from j1 limit 10);", parser_dialect: MySqlDialect {}, unparser_dialect: UnparserMySqlDialect {}, - expected: @"SELECT `j1`.`j1_id` FROM (SELECT `j1`.`j1_id` FROM `j1` LIMIT 10) AS `derived_limit`", + expected: @"SELECT `derived_limit`.`j1_id` FROM (SELECT `j1`.`j1_id` FROM `j1` LIMIT 10) AS `derived_limit`", ); Ok(()) } @@ -1614,7 +1695,7 @@ fn test_table_scan_pushdown() -> Result<()> { plan_to_sql(&query_from_table_scan_with_two_projections)?; assert_snapshot!( query_from_table_scan_with_two_projections, - @"SELECT t1.id, t1.age FROM (SELECT t1.id, t1.age FROM t1)" + @"SELECT id, age FROM (SELECT t1.id, t1.age FROM t1)" ); let table_scan_with_filter = table_scan_with_filters( @@ -1791,7 +1872,7 @@ fn test_sort_with_scalar_fn_and_push_down_fetch() -> Result<()> { let sql = plan_to_sql(&plan)?; assert_snapshot!( sql, - @"SELECT t1.search_phrase FROM (SELECT t1.search_phrase, t1.event_time FROM t1 WHERE (t1.search_phrase <> '') ORDER BY substr(t1.event_time, 1, 5) ASC NULLS FIRST LIMIT 10)" + @"SELECT search_phrase FROM (SELECT t1.search_phrase, t1.event_time FROM t1 WHERE (t1.search_phrase <> '') ORDER BY substr(t1.event_time, 1, 5) ASC NULLS FIRST LIMIT 10)" ); Ok(()) } diff --git a/datafusion/sqllogictest/README.md b/datafusion/sqllogictest/README.md index 76c28b5ecd2ec..ca00fff7f6544 100644 --- a/datafusion/sqllogictest/README.md +++ b/datafusion/sqllogictest/README.md @@ -349,7 +349,7 @@ export RUST_MIN_STACK=30485760; PG_COMPAT=true INCLUDE_SQLITE=true cargo test --features=postgres --test sqllogictests ``` -To update the sqllite expected answers use the `datafusion/sqllogictest/regenerate_sqlite_files.sh` script. +To update the sqlite expected answers use the `datafusion/sqllogictest/regenerate_sqlite_files.sh` script. Note this must be run with an empty postgres instance. For example diff --git a/datafusion/sqllogictest/test_files/array/array_position.slt b/datafusion/sqllogictest/test_files/array/array_position.slt index e3dd830dfb77a..8fe1826619431 100644 --- a/datafusion/sqllogictest/test_files/array/array_position.slt +++ b/datafusion/sqllogictest/test_files/array/array_position.slt @@ -282,6 +282,17 @@ select array_position([1, 2, 3], 3, 4), array_position([1], 1, 2); ---- NULL NULL +query error start_from out of bounds: -9223372036854775808, expected -9223372036854775807 to 9223372036854775807 +select array_position([1], 1, -9223372036854775808); + +query error start_from out of bounds: -9223372036854775808, expected -9223372036854775807 to 9223372036854775807 +select array_position([1], 1, start_from) +from (values (-9223372036854775808)) as t(start_from); + +query error start_from out of bounds: -9223372036854775808, expected -9223372036854775807 to 9223372036854775807 +select array_position([1], needle, start_from) +from (values (1, -9223372036854775808)) as t(needle, start_from); + # array_position with empty array in various contexts query II select array_position(arrow_cast(make_array(), 'List(Int64)'), 1), array_position(arrow_cast(make_array(), 'LargeList(Int64)'), 1); diff --git a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt index 425d8ac16eaee..01a19454e9a80 100644 --- a/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt +++ b/datafusion/sqllogictest/test_files/cast_extension_type_metadata.slt @@ -45,5 +45,55 @@ FROM ( ---- 00010203040506070809000102030506 arrow.uuid -statement error DataFusion error: Optimizer rule 'simplify_expressions' failed[\s\S]*TryCast from FixedSizeBinary\(16\) to FixedSizeBinary\(16\)<\{"ARROW:extension:name": "arrow\.uuid"\}> is not supported -SELECT TRY_CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID); +# TRY_CAST to extension type should also preserve extension metadata +query ?T +SELECT + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + arrow_metadata( + TRY_CAST( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + AS UUID + ), + 'ARROW:extension:name' + ); +---- +00010203040506070809000102030506 arrow.uuid + +# TRY_CAST to UUID from a subquery +query ?T +SELECT + TRY_CAST(raw AS UUID), + arrow_metadata(TRY_CAST(raw AS UUID), 'ARROW:extension:name') +FROM ( + VALUES ( + arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') + ) +) AS uuids(raw); +---- +00010203040506070809000102030506 arrow.uuid + +# arrow_cast from UUID to same underlying type (FixedSizeBinary(16)) strips +# extension metadata (type-only cast semantics) +query ?T +SELECT + arrow_cast(uuid_val, 'FixedSizeBinary(16)'), + arrow_metadata(arrow_cast(uuid_val, 'FixedSizeBinary(16)'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL + +# arrow_cast to a different type strips extension metadata (type-only cast semantics) +query ?T +SELECT + arrow_cast(uuid_val, 'Binary'), + arrow_metadata(arrow_cast(uuid_val, 'Binary'), 'ARROW:extension:name') +FROM ( + SELECT CAST(arrow_cast(X'00010203040506070809000102030506', 'FixedSizeBinary(16)') AS UUID) AS uuid_val +); +---- +00010203040506070809000102030506 NULL diff --git a/datafusion/sqllogictest/test_files/mark_join_matrix.slt b/datafusion/sqllogictest/test_files/mark_join_matrix.slt new file mode 100644 index 0000000000000..a17227df48e4b --- /dev/null +++ b/datafusion/sqllogictest/test_files/mark_join_matrix.slt @@ -0,0 +1,201 @@ +# 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. + +# Mark-join correctness across a config matrix. A mark join adds a boolean `mark` +# per left row and comes from EXISTS/IN/NOT EXISTS/NOT IN inside a disjunction +# (`WHERE OR EXISTS(..)`; see mark_join in decorrelate_predicate_subquery.rs). +# For an equijoin mark, prefer_hash_join toggles the operator: false -> +# SortMergeJoinExec LeftMark, true -> HashJoinExec RightMark (inputs swapped), so +# sweeping {true,false} x batch_size {1,2,100,8192} cross-checks both directions +# and must agree row-for-row. Range mark joins run on NestedLoopJoin regardless +# (Part 2, batch_size only). SMJ needs target_partitions>1 and repartition_joins, +# set below (not swept). Matrix rules: no EXPLAIN, no in-file SET of a swept knob, +# rowsort every multi-row query. + +# configMatrix: datafusion.optimizer.prefer_hash_join=true,false +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +# ------------------------------------------------------------------ +# Fixtures: duplicate and NULL keys on the subquery side, a NULL key and an +# unmatched key on the outer side. +# ------------------------------------------------------------------ +statement ok +CREATE TABLE mk_l(k INT, v INT) AS VALUES + (1, 10), (2, 20), (3, 30), (4, 40), (NULL, 50); + +statement ok +CREATE TABLE mk_r(k INT) AS VALUES (1), (2), (2), (NULL); + +# mk_r without the NULL, for NOT IN cases that would otherwise be swallowed by +# three-valued logic. +statement ok +CREATE TABLE mk_r_nn(k INT) AS VALUES (1), (2), (2); + +statement ok +CREATE TABLE mk_empty(k INT); + +# ================================================================== +# Part 1: Equijoin mark joins (LeftMark on SMJ vs RightMark on HashJoin) +# ================================================================== + +# EXISTS in a disjunction with a sometimes-true predicate. mark(EXISTS k in +# {1,2}) is true for k=1,2; the predicate v>35 is true for k=4 and k=NULL. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +1 10 +2 20 +4 40 +NULL 50 + +# NOT EXISTS in a disjunction: mark is negated, true for k not in {1,2}. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR NOT EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +3 30 +4 40 +NULL 50 + +# Predicate never true (no negative k), so the result isolates the mark: the +# EXISTS rows k in {1,2}. Exercises the mark column with the OR contributing +# nothing. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +1 10 +2 20 + +# Same, negated: isolates NOT EXISTS. The NULL-keyed left row never matches, so +# it is kept. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR NOT EXISTS (SELECT 1 FROM mk_r r WHERE l.k = r.k); +---- +3 30 +4 40 +NULL 50 + +# IN in a disjunction: same matches as EXISTS here; the NULL in the subquery adds +# no true values. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR l.k IN (SELECT r.k FROM mk_r r); +---- +1 10 +2 20 +4 40 +NULL 50 + +# NOT IN over a NULL-free subquery: mark(NOT IN) is true for k not in {1,2}. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r_nn r); +---- +3 30 +4 40 +NULL 50 + +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR l.k NOT IN (SELECT r.k FROM mk_r r); +---- +4 40 +NULL 50 + +# Empty subquery: EXISTS is always false, so the result is just the predicate. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 35 OR EXISTS (SELECT 1 FROM mk_empty r WHERE l.k = r.k); +---- +4 40 +NULL 50 + +# Empty subquery, negated: NOT EXISTS is always true, so every row survives. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR NOT EXISTS (SELECT 1 FROM mk_empty r WHERE l.k = r.k); +---- +1 10 +2 20 +3 30 +4 40 +NULL 50 + +# Multiple equi keys in the mark correlation. +statement ok +CREATE TABLE mk2_l(k1 INT, k2 INT, v INT) AS VALUES + (1, 1, 10), (1, 2, 20), (2, 2, 30), (3, 3, 40); + +statement ok +CREATE TABLE mk2_r(k1 INT, k2 INT) AS VALUES (1, 1), (2, 2), (1, 9); + +query III rowsort +SELECT l.k1, l.k2, l.v FROM mk2_l l +WHERE l.v > 35 OR EXISTS (SELECT 1 FROM mk2_r r WHERE l.k1 = r.k1 AND l.k2 = r.k2); +---- +1 1 10 +2 2 30 +3 3 40 + +# Cross-table filter in the correlation (regression #21197): unmatched mark rows +# produce null right indices that must not corrupt non-nullable left columns. +# EXISTS holds for the k=2 rows (partner (2,99) differs in d); k=1's partner +# equals its d, so false. +statement ok +CREATE TABLE mkf_l(k INT, d INT) AS VALUES (1, 10), (2, 20), (2, 25), (3, 30); + +statement ok +CREATE TABLE mkf_r(k INT, d INT) AS VALUES (1, 10), (2, 20), (2, 99); + +query II rowsort +SELECT l.k, l.d FROM mkf_l l +WHERE l.d < 0 OR EXISTS (SELECT 1 FROM mkf_r r WHERE l.k = r.k AND r.d <> l.d); +---- +2 20 +2 25 + +# ================================================================== +# Part 2: Range mark joins (RightMark on NestedLoopJoin) +# ================================================================== +# No equi key, so these run on NestedLoopJoin in both combinations; only +# batch_size varies. mk_r non-null keys = {1,2,2}. + +# EXISTS l.k > r.k is true for k>1, i.e. k in {2,3,4}. Predicate never true. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.v > 100 OR EXISTS (SELECT 1 FROM mk_r r WHERE l.k > r.k); +---- +2 20 +3 30 +4 40 + +# NOT EXISTS of the same range: true for k=1 (no smaller r) and the NULL key. +query II rowsort +SELECT l.k, l.v FROM mk_l l +WHERE l.k < 0 OR NOT EXISTS (SELECT 1 FROM mk_r r WHERE l.k > r.k); +---- +1 10 +NULL 50 diff --git a/datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt b/datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt new file mode 100644 index 0000000000000..db27c9fc8bb60 --- /dev/null +++ b/datafusion/sqllogictest/test_files/piecewise_merge_join_batches.slt @@ -0,0 +1,103 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Plan proof for piecewise_merge_join_matrix.slt. Not a matrix file: it fixes +# batch_size=2 and enable_piecewise_merge_join=true so it can EXPLAIN what the +# matrix relies on but cannot check itself -- the streamed side is a range +# LazyMemoryExec with batch_size=2 (so range(3,8), 5 rows, arrives as 3 batches) +# feeding a PiecewiseMergeJoin, not a fallback. A VALUES source would not show +# batch_size in the plan. + +statement ok +set datafusion.optimizer.enable_piecewise_merge_join = true; + +statement ok +set datafusion.execution.batch_size = 2; + +statement ok +CREATE TABLE pb_l(id INT, v INT) AS + SELECT CAST(value AS INT) AS id, CAST(value AS INT) AS v FROM range(1, 11); + +# Existence (LeftSemi): range streamed side (batch_size=2) -> LeftSemi PWMJ. +query TT +EXPLAIN SELECT count(*) FROM pb_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----Projection: +04)------LeftSemi Join: Filter: CAST(l.v AS Int64) > __correlated_sq_1.value +05)--------SubqueryAlias: l +06)----------TableScan: pb_l projection=[v] +07)--------SubqueryAlias: __correlated_sq_1 +08)----------SubqueryAlias: r +09)------------TableScan: range() projection=[value] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)] +02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +03)----CoalescePartitionsExec +04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +05)--------ProjectionExec: expr=[] +06)----------PiecewiseMergeJoin: operator=Gt, join_type=LeftSemi, on=(CAST(v AS Int64) > value) +07)------------SortPreservingMergeExec: [CAST(v@0 AS Int64) ASC] +08)--------------SortExec: expr=[CAST(v@0 AS Int64) ASC], preserve_partitioning=[true] +09)----------------DataSourceExec: partitions=4, partition_sizes=[2, 1, 1, 1] +10)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +11)--------------LazyMemoryExec: partitions=1, batch_generators=[range: start=3, end=8, batch_size=2] + +# Same result the matrix asserts (v > min(3) -> {4..10} = 7). +query I +SELECT count(*) FROM pb_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +7 + +# Classic Inner: same range streamed side -> Inner PWMJ. +query TT +EXPLAIN SELECT count(*) FROM pb_l l JOIN range(3, 8) r ON l.v < r.value; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----Projection: +04)------Inner Join: Filter: CAST(l.v AS Int64) < r.value +05)--------SubqueryAlias: l +06)----------TableScan: pb_l projection=[v] +07)--------SubqueryAlias: r +08)----------TableScan: range() projection=[value] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)] +02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +03)----CoalescePartitionsExec +04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +05)--------ProjectionExec: expr=[] +06)----------PiecewiseMergeJoin: operator=Lt, join_type=Inner, on=(CAST(v AS Int64) < value) +07)------------SortPreservingMergeExec: [CAST(v@0 AS Int64) DESC] +08)--------------SortExec: expr=[CAST(v@0 AS Int64) DESC], preserve_partitioning=[true] +09)----------------DataSourceExec: partitions=4, partition_sizes=[2, 1, 1, 1] +10)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +11)--------------LazyMemoryExec: partitions=1, batch_generators=[range: start=3, end=8, batch_size=2] + +query I +SELECT count(*) FROM pb_l l JOIN range(3, 8) r ON l.v < r.value; +---- +20 + +statement ok +RESET datafusion.execution.batch_size; + +statement ok +RESET datafusion.optimizer.enable_piecewise_merge_join; diff --git a/datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt b/datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt new file mode 100644 index 0000000000000..13aeeddfc2afa --- /dev/null +++ b/datafusion/sqllogictest/test_files/piecewise_merge_join_matrix.slt @@ -0,0 +1,661 @@ +# 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. + +# PiecewiseMergeJoin correctness across a config matrix. A range join with no +# equi key runs on PiecewiseMergeJoin when enable_piecewise_merge_join=true and +# on NestedLoopJoin when false, so sweeping {true,false} x batch_size +# {1,2,100,8192} requires the two to agree row-for-row. PWMJ covers +# Inner/Left/Right/Full (Part 1) and LeftSemi/LeftAnti (Part 2); explicit RIGHT +# SEMI/ANTI JOIN and Mark (OR EXISTS) stay on NLJ. Matrix rules: no EXPLAIN, no +# in-file SET of a swept knob, rowsort every multi-row query. + +# configMatrix: datafusion.optimizer.enable_piecewise_merge_join=true,false +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +# ================================================================== +# Part 1: Classic range joins (Inner / Left / Right / Full) +# ================================================================== +# cj_l.lv = {1,3,5,NULL}, cj_r.rv = {3,5,NULL}. A NULL key matches nothing, so it +# appears only as an unmatched row on an outer side. +statement ok +CREATE TABLE cj_l(lid INT, lv INT); + +statement ok +INSERT INTO cj_l VALUES (1, 1), (2, 3), (3, 5), (4, NULL); + +statement ok +CREATE TABLE cj_r(rid INT, rv INT); + +statement ok +INSERT INTO cj_r VALUES (1, 3), (2, 5), (3, NULL); + +# INNER, all four operators. `<=`/`>=` differ from `<`/`>` through the equal +# values (lv=3=rv, lv=5=rv). +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv <= r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 1 3 +2 3 2 5 +3 5 2 5 + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv > r.rv; +---- +3 5 1 3 + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_r r ON l.lv >= r.rv; +---- +2 3 1 3 +3 5 1 3 +3 5 2 5 + +# LEFT: the INNER `<` rows plus every unmatched left row (lv=5 and lv=NULL) with +# a NULL right side. +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l LEFT JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 +3 5 NULL NULL +4 NULL NULL NULL + +# RIGHT: INNER rows plus unmatched right rows (NULL left). rid=3 has a NULL key +# and must still be emitted (regression: streamed-side NULLs were dropped). +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l RIGHT JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 +NULL NULL 3 NULL + +# RIGHT with a different operator flips the buffered-side sort direction; the +# NULL-keyed right row rid=3 is still emitted unmatched. +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l RIGHT JOIN cj_r r ON l.lv >= r.rv; +---- +2 3 1 3 +3 5 1 3 +3 5 2 5 +NULL NULL 3 NULL + +# FULL: the INNER rows plus unmatched rows from both sides, including a NULL key +# on each side. +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l FULL JOIN cj_r r ON l.lv < r.rv; +---- +1 1 1 3 +1 1 2 5 +2 3 2 5 +3 5 NULL NULL +4 NULL NULL NULL +NULL NULL 3 NULL + +# Empty right side: LEFT keeps every left row with NULLs, INNER is empty. +statement ok +CREATE TABLE cj_empty(rid INT, rv INT); + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l JOIN cj_empty r ON l.lv < r.rv; +---- + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_l l LEFT JOIN cj_empty r ON l.lv < r.rv; +---- +1 1 NULL NULL +2 3 NULL NULL +3 5 NULL NULL +4 NULL NULL NULL + +# Empty left side: RIGHT keeps every right row with NULLs, including rid=3's NULL +# key. +query IIII rowsort +SELECT l.rid, l.rv, r.rid, r.rv FROM cj_empty l RIGHT JOIN cj_r r ON l.rv < r.rv; +---- +NULL NULL 1 3 +NULL NULL 2 5 +NULL NULL 3 NULL + +# Duplicate keys on both sides: each qualifying left row must pair with the whole +# matching right suffix. lv=1 (twice) is below rv=2 (twice) -> 4 rows; lv=3 has +# no larger right value. +statement ok +CREATE TABLE cj_dup_l(lid INT, lv INT); + +statement ok +INSERT INTO cj_dup_l VALUES (1, 1), (2, 1), (3, 3); + +statement ok +CREATE TABLE cj_dup_r(rid INT, rv INT); + +statement ok +INSERT INTO cj_dup_r VALUES (1, 2), (2, 2); + +query IIII rowsort +SELECT l.lid, l.lv, r.rid, r.rv FROM cj_dup_l l JOIN cj_dup_r r ON l.lv < r.rv; +---- +1 1 1 2 +1 1 2 2 +2 1 1 2 +2 1 2 2 + +# ================================================================== +# Part 2: Existence joins (LeftSemi / LeftAnti; Right/Mark fall back to NLJ) +# ================================================================== +# EXISTS -> LeftSemi, NOT EXISTS -> LeftAnti (the only existence joins PWMJ +# implements). Explicit RIGHT SEMI/ANTI JOIN and Mark (OR EXISTS) stay on NLJ; +# covered at the end. + +# ------------------------------------------------------------------ +# Fixtures +# ------------------------------------------------------------------ +statement ok +CREATE TABLE ej_l(id INT, v INT); + +statement ok +INSERT INTO ej_l VALUES (1, 5), (2, 4), (3, 2), (4, 1); + +statement ok +CREATE TABLE ej_r(v INT); + +statement ok +INSERT INTO ej_r VALUES (2), (3), (4); + +# ------------------------------------------------------------------ +# Basic: every operator, both LeftSemi (EXISTS) and LeftAnti (NOT EXISTS) +# ------------------------------------------------------------------ +# ej_l.v = {5,4,2,1}, ej_r.v = {2,3,4}. + +# `<` : 2<3 and 1<2 match; 5 and 4 exceed every right value. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v < r.v); +---- +3 +4 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v < r.v); +---- +1 +2 + +# `<=` : additionally admits v=4 through the equal right value 4. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v <= r.v); +---- +2 +3 +4 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v <= r.v); +---- +1 + +# `>` : 5 and 4 exceed some right value; 2 and 1 do not. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +1 +2 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +3 +4 + +# `>=` : additionally admits v=2 through the equal right value 2. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v >= r.v); +---- +1 +2 +3 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v >= r.v); +---- +4 + +# ------------------------------------------------------------------ +# Correlation written inner-column-first, and with an expression +# ------------------------------------------------------------------ +# `r.v < l.v` is the same join as `l.v > r.v`; the planner flips the operator so +# the marked (left) side stays buffered. Same rows as `l.v > r.v` above. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE r.v < l.v); +---- +1 +2 + +# An expression on the streamed side: `l.v < r.v + 1` is `l.v <= r.v` over ints. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v < r.v + 1); +---- +2 +3 +4 + +# ------------------------------------------------------------------ +# NULL semantics on the buffered (left) key +# ------------------------------------------------------------------ +# A NULL key satisfies no comparison, so it is excluded from EXISTS and kept by +# NOT EXISTS. ej_ln.v = {NULL,3,NULL,1}, right non-null key = {2}. +statement ok +CREATE TABLE ej_ln(id INT, v INT); + +statement ok +INSERT INTO ej_ln VALUES (1, NULL), (2, 3), (3, NULL), (4, 1); + +statement ok +CREATE TABLE ej_rn(v INT); + +statement ok +INSERT INTO ej_rn VALUES (2), (NULL); + +# `<` : only v=1 is below 2. Both NULL-keyed left rows stay in NOT EXISTS. +query I rowsort +SELECT l.id FROM ej_ln l WHERE EXISTS (SELECT 1 FROM ej_rn r WHERE l.v < r.v); +---- +4 + +query I rowsort +SELECT l.id FROM ej_ln l WHERE NOT EXISTS (SELECT 1 FROM ej_rn r WHERE l.v < r.v); +---- +1 +2 +3 + +# `>` : only v=3 exceeds 2. +query I rowsort +SELECT l.id FROM ej_ln l WHERE EXISTS (SELECT 1 FROM ej_rn r WHERE l.v > r.v); +---- +2 + +query I rowsort +SELECT l.id FROM ej_ln l WHERE NOT EXISTS (SELECT 1 FROM ej_rn r WHERE l.v > r.v); +---- +1 +3 +4 + +# ------------------------------------------------------------------ +# Empty inputs +# ------------------------------------------------------------------ +statement ok +CREATE TABLE ej_empty(v INT); + +# Empty streamed (right) side: no key can match, so EXISTS is empty and NOT +# EXISTS keeps every left row. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_empty r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_empty r WHERE l.v > r.v); +---- +1 +2 +3 +4 + +statement ok +CREATE TABLE ej_lempty(id INT, v INT); + +# Empty buffered (left) side: nothing to emit either way. +query I rowsort +SELECT l.id FROM ej_lempty l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_lempty l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- + +# ------------------------------------------------------------------ +# All-NULL sides (distinct early-exit triggers in the existence path) +# ------------------------------------------------------------------ +# An all-NULL buffered side saturates the watermark before the first poll, so no +# streamed batch is read; EXISTS is still empty. +statement ok +CREATE TABLE ej_l_allnull(id INT, v INT); + +statement ok +INSERT INTO ej_l_allnull VALUES (1, NULL), (2, NULL); + +query I rowsort +SELECT l.id FROM ej_l_allnull l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_l_allnull l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +1 +2 + +# An all-NULL streamed side never lowers the watermark, so it behaves like an +# empty streamed side. +statement ok +CREATE TABLE ej_r_allnull(v INT); + +statement ok +INSERT INTO ej_r_allnull VALUES (NULL), (NULL); + +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r_allnull r WHERE l.v > r.v); +---- + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r_allnull r WHERE l.v > r.v); +---- +1 +2 +3 +4 + +# ------------------------------------------------------------------ +# Duplicate buffered keys straddling the match boundary +# ------------------------------------------------------------------ +# The first matching buffered row is found by binary search; it must return the +# FIRST index of a run of equal keys or earlier duplicates vanish from EXISTS. +# ej_dup_l.v = {5,5,3,3,1}; the deciding streamed key is 3, so `<=`/`>=` land +# inside the run of 3s. +statement ok +CREATE TABLE ej_dup_l(id INT, v INT); + +statement ok +INSERT INTO ej_dup_l VALUES (1, 5), (2, 5), (3, 3), (4, 3), (5, 1); + +statement ok +CREATE TABLE ej_dup_r(v INT); + +statement ok +INSERT INTO ej_dup_r VALUES (3); + +# `<` : only v=1 is below 3. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v < r.v); +---- +5 + +# `<=` : both v=3 rows must appear. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v <= r.v); +---- +3 +4 +5 + +# `>` : only the two v=5 rows exceed 3. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v > r.v); +---- +1 +2 + +# `>=` : the boundary is inside the run of 3s from the other direction. +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v >= r.v); +---- +1 +2 +3 +4 + +query I rowsort +SELECT l.id FROM ej_dup_l l WHERE NOT EXISTS (SELECT 1 FROM ej_dup_r r WHERE l.v >= r.v); +---- +5 + +# ------------------------------------------------------------------ +# Subquery filter -> repartitioned streamed side (multi-partition final pass) +# ------------------------------------------------------------------ +# A predicate on the inner table is pushed to the scan and lets the streamed +# side repartition, so the final existence pass is coordinated across streamed +# partitions. `r.v > 0` keeps all of {2,3,4}, so the rows match `l.v > r.v`. +query I rowsort +SELECT l.id FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v AND r.v > 0); +---- +1 +2 + +query I rowsort +SELECT l.id FROM ej_l l WHERE NOT EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v AND r.v > 0); +---- +3 +4 + +# ------------------------------------------------------------------ +# Type coverage: the existence path shares the join comparator +# ------------------------------------------------------------------ + +# Date32 key. +statement ok +CREATE TABLE ej_dl(id INT, d DATE); + +statement ok +INSERT INTO ej_dl VALUES (1, DATE '2022-04-23'), (2, DATE '2022-04-28'), (3, DATE '2022-04-18'); + +statement ok +CREATE TABLE ej_dr(d DATE); + +statement ok +INSERT INTO ej_dr VALUES (DATE '2022-04-20'), (DATE '2022-04-26'); + +query I rowsort +SELECT l.id FROM ej_dl l WHERE EXISTS (SELECT 1 FROM ej_dr r WHERE l.d > r.d); +---- +1 +2 + +query I rowsort +SELECT l.id FROM ej_dl l WHERE NOT EXISTS (SELECT 1 FROM ej_dr r WHERE l.d > r.d); +---- +3 + +# Float key with negative zero: -0.0 and +0.0 compare equal in SQL, so +# `-0.0 < 0.0` is false and `-0.0 <= 0.0` is true. The comparator normalizes the +# sign of zero before comparing. +statement ok +CREATE TABLE ej_fl(id INT, v DOUBLE); + +statement ok +INSERT INTO ej_fl VALUES (1, -0.0), (2, 2.5); + +statement ok +CREATE TABLE ej_fr(v DOUBLE); + +statement ok +INSERT INTO ej_fr VALUES (0.0); + +query I rowsort +SELECT l.id FROM ej_fl l WHERE EXISTS (SELECT 1 FROM ej_fr r WHERE l.v < r.v); +---- + +query I rowsort +SELECT l.id FROM ej_fl l WHERE EXISTS (SELECT 1 FROM ej_fr r WHERE l.v <= r.v); +---- +1 + +# String key. +statement ok +CREATE TABLE ej_sl(id INT, s VARCHAR); + +statement ok +INSERT INTO ej_sl VALUES (1, 'apple'), (2, 'cherry'), (3, 'mango'); + +statement ok +CREATE TABLE ej_sr(s VARCHAR); + +statement ok +INSERT INTO ej_sr VALUES ('banana'), ('lemon'); + +query I rowsort +SELECT l.id FROM ej_sl l WHERE EXISTS (SELECT 1 FROM ej_sr r WHERE l.s > r.s); +---- +2 +3 + +query I rowsort +SELECT l.id FROM ej_sl l WHERE NOT EXISTS (SELECT 1 FROM ej_sr r WHERE l.s > r.s); +---- +1 + +# Dictionary-encoded key: no typed arrow min/max kernel, so the extreme key per +# streamed batch is chosen through the generic ScalarValue path. +statement ok +CREATE TABLE ej_dict_l AS + SELECT column1 AS id, arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS v + FROM (VALUES (1, 'a'), (2, 'c'), (3, 'e'), (4, NULL)); + +statement ok +CREATE TABLE ej_dict_r AS + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS v FROM (VALUES ('c')); + +query I rowsort +SELECT l.id FROM ej_dict_l l WHERE EXISTS (SELECT 1 FROM ej_dict_r r WHERE l.v > r.v); +---- +3 + +query I rowsort +SELECT l.id FROM ej_dict_l l WHERE NOT EXISTS (SELECT 1 FROM ej_dict_r r WHERE l.v > r.v); +---- +1 +2 +4 + +query I rowsort +SELECT l.id FROM ej_dict_l l WHERE EXISTS (SELECT 1 FROM ej_dict_r r WHERE l.v <= r.v); +---- +1 +2 + +# ------------------------------------------------------------------ +# Multi-batch stress: streamed side from range() (fragments by batch_size) +# ------------------------------------------------------------------ +# range() is a LazyMemoryExec that emits ceil(n / batch_size) batches, so +# batch_size genuinely fragments the streamed side (unlike a VALUES table, whose +# split is incidental). Plan proof in piecewise_merge_join_batches.slt. Buffered +# v = 1..10; streamed range(3,8) = {3,4,5,6,7}; counts are matrix-invariant. +statement ok +CREATE TABLE ej_range_l(id INT, v INT) AS + SELECT CAST(value AS INT) AS id, CAST(value AS INT) AS v FROM range(1, 11); + +# Existence, `>`: v > min(3) -> {4..10} = 7. +query I +SELECT count(*) FROM ej_range_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +7 + +query I +SELECT count(*) FROM ej_range_l l WHERE NOT EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v > r.value); +---- +3 + +# `<`: v < max(7) -> {1..6} = 6. +query I +SELECT count(*) FROM ej_range_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v < r.value); +---- +6 + +query I +SELECT count(*) FROM ej_range_l l WHERE NOT EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v < r.value); +---- +4 + +# `>=` with a pushed-down inner filter that keeps every row, so the streamed side +# also repartitions. v >= min(3) -> {3..10} = 8. +query I +SELECT count(*) FROM ej_range_l l WHERE EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v >= r.value AND r.value > 0); +---- +8 + +query I +SELECT count(*) FROM ej_range_l l WHERE NOT EXISTS (SELECT 1 FROM range(3, 8) r WHERE l.v >= r.value AND r.value > 0); +---- +2 + +# Classic range join over the same range() streamed side. INNER l.v < r.value: +# sum over v of |{r : r > v}| = 5+5+4+3+2+1 = 20. +query I +SELECT count(*) FROM ej_range_l l JOIN range(3, 8) r ON l.v < r.value; +---- +20 + +# LEFT adds the unmatched left rows (v >= 7, which is v in {7,8,9,10}): 20 + 4 = 24. +query I +SELECT count(*) FROM ej_range_l l LEFT JOIN range(3, 8) r ON l.v < r.value; +---- +24 + +# Existence result feeding an aggregate: EXISTS l.v>r.v over the small fixture is +# {1,2}, so the count is 2. +query I +SELECT count(*) FROM ej_l l WHERE EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +2 + +# ------------------------------------------------------------------ +# Other existence joins: RightSemi / RightAnti / Mark +# ------------------------------------------------------------------ +# PWMJ rejects these, so both matrix combinations run them on NestedLoopJoin; +# only batch_size varies. Included so the file covers every existence join type. +statement ok +CREATE TABLE ej_rs_l(v INT); + +statement ok +INSERT INTO ej_rs_l VALUES (3), (6); + +statement ok +CREATE TABLE ej_rs_r(id INT, v INT); + +statement ok +INSERT INTO ej_rs_r VALUES (1, 1), (2, 4), (3, 7); + +# RIGHT SEMI: keep right rows that have some larger left value. left={3,6}: +# right 1 (l=3,6), right 4 (l=6) qualify; right 7 has none. +query I rowsort +SELECT r.id FROM ej_rs_l l RIGHT SEMI JOIN ej_rs_r r ON l.v > r.v; +---- +1 +2 + +# RIGHT ANTI: the complement -> only right 7. +query I rowsort +SELECT r.id FROM ej_rs_l l RIGHT ANTI JOIN ej_rs_r r ON l.v > r.v; +---- +3 + +# Mark join: EXISTS used inside a disjunction produces a LeftMark. `l.v > r.v` +# over the fixture yields {1,2}; `l.id = 4` adds id 4. +query I rowsort +SELECT l.id FROM ej_l l WHERE l.id = 4 OR EXISTS (SELECT 1 FROM ej_r r WHERE l.v > r.v); +---- +1 +2 +4 diff --git a/datafusion/sqllogictest/test_files/projection.slt b/datafusion/sqllogictest/test_files/projection.slt index e18114bc51ca8..bac94670f4ffd 100644 --- a/datafusion/sqllogictest/test_files/projection.slt +++ b/datafusion/sqllogictest/test_files/projection.slt @@ -217,6 +217,74 @@ SELECT column1 as a from (values (1), (2)) f where f.column1 = 2; ---- 2 +# Regression: preserve deep anonymous nested projections with one and default passes. +statement ok +CREATE TABLE nested_projection(i INT); + +statement ok +INSERT INTO nested_projection VALUES (3), (4), (5); + +statement ok +SET datafusion.optimizer.max_passes = 1; + +query I rowsort +SELECT i +FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM nested_projection + ) + ) + ) + ) + ) +) +---- +10 +11 +9 + +statement ok +RESET datafusion.optimizer.max_passes; + +query I rowsort +SELECT i +FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM ( + SELECT i + 1 AS i + FROM nested_projection + ) + ) + ) + ) + ) +) +---- +10 +11 +9 + +statement ok +DROP TABLE nested_projection; + # clean data statement ok DROP TABLE aggregate_simple; diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index 3c7b6f4cb1127..1f9176f7137bd 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -2144,6 +2144,83 @@ true true true +##################### +# Section: leaf expression extraction does not duplicate volatile expressions +# +# Regression test for #24678. `extract_leaf_expressions` / `push_down_leaf_projections` +# build an extraction projection by *merging* into the input projection, which +# inlines the definition of every column the extracted expression references. +# When that definition is volatile the inlined copy is an independent +# evaluation, so `random()` is drawn twice and the extracted value no longer +# matches the struct it was derived from. +# +# Before the fix the plan below read +# Projection: named_struct(Utf8("a"), random()) AS s, random() AS field +# and `field` differed from `s['a']` on every row. This is the same invariant +# `FileScanConfig::try_swapping_with_projection` already enforces for the +# physical pushdown path (#23220). +##################### + +statement ok +set datafusion.explain.logical_plan_only = true; + +# `random()` must appear once; `field` reads the struct materialized below it. +query TT +EXPLAIN SELECT s, s['a'] AS field +FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 3)); +---- +logical_plan +01)Projection: s, get_field(s, Utf8("a")) AS field +02)--Projection: named_struct(Utf8("a"), random()) AS s +03)----TableScan: generate_series() projection=[] + +# Same for an extraction out of a Filter: the predicate must test the struct +# that is returned, not a second draw. +query TT +EXPLAIN SELECT s +FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 3)) +WHERE s['a'] > 0.5; +---- +logical_plan +01)Filter: get_field(s, Utf8("a")) > Float64(0.5) +02)--Projection: named_struct(Utf8("a"), random()) AS s +03)----TableScan: generate_series() projection=[] + +statement ok +set datafusion.explain.logical_plan_only = false; + +# Every row's `field` is the `a` field of that row's own struct. +query B +SELECT bool_and(field = s['a']) +FROM ( + SELECT s, s['a'] AS field + FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) +); +---- +true + +# Same invariant when the extraction is routed into one side of a join. +query B +SELECT bool_and(field = s['a']) +FROM ( + SELECT l.s, l.s['a'] AS field + FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) AS l + INNER JOIN (SELECT 1 AS k) AS r ON true +); +---- +true + +# Every surviving row satisfies the predicate it was filtered on. +query B +SELECT bool_and(s['a'] > 0.5) +FROM ( + SELECT s + FROM (SELECT named_struct('a', random()) AS s FROM generate_series(1, 1000)) + WHERE s['a'] > 0.5 +); +---- +true + ##################### # Section: expensive expressions are not re-inlined by projection pushdown # diff --git a/datafusion/sqllogictest/test_files/run_end_encoded.slt b/datafusion/sqllogictest/test_files/run_end_encoded.slt index b5909bc7c430d..094f6ede1b294 100644 --- a/datafusion/sqllogictest/test_files/run_end_encoded.slt +++ b/datafusion/sqllogictest/test_files/run_end_encoded.slt @@ -85,3 +85,65 @@ false false true true true true true true + +# date_bin's signature is a plain Timestamp, not RunEndEncoded, so this only +# plans if coerced_from() unwraps the REE value type the same way it does for +# Dictionary (see test_coerced_from_run_end_encoded) +statement ok +CREATE TABLE ree_timestamps AS +SELECT + arrow_cast( + arrow_cast(ts, 'Timestamp(Nanosecond, None)'), + 'RunEndEncoded("run_ends": non-null Int32, "values": Timestamp(Nanosecond, None))' + ) AS ts +FROM (VALUES + ('2023-12-04T00:00:00'), + ('2023-12-04T00:05:00'), + ('2023-12-04T00:35:00') +) AS t(ts); + +query P rowsort +SELECT date_bin('30 minutes', ts) FROM ree_timestamps; +---- +2023-12-04T00:00:00 +2023-12-04T00:00:00 +2023-12-04T00:30:00 + +# RANGE window frame ordered by an REE-encoded column: extract_window_frame_target_type() +# unwraps RunEndEncoded the same way it does for Dictionary, so a frame that +# doesn't need offset arithmetic. +query II +SELECT + temperature, + SUM(temperature) OVER ( + ORDER BY arrow_cast( + temperature, + 'RunEndEncoded("run_ends": non-null Int32, "values": Int64)' + ) + RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW + ) AS windowed_sum +FROM sensor_readings +ORDER BY temperature; +---- +20 20 +22 42 +23 65 +24 89 + +# An offset RANGE bound (`k PRECEDING`/`k FOLLOWING`, unlike CURRENT ROW or +# UNBOUNDED above) is computed by adding/subtracting the offset from the +# current row's value, and arrow's numeric kernels don't support arithmetic +# on a still-RunEndEncoded-wrapped value (tracked at +# https://github.com/apache/arrow-rs/issues/10891). +query error RANGE with offset PRECEDING/FOLLOWING is not supported for ORDER BY type +SELECT + temperature, + SUM(temperature) OVER ( + ORDER BY arrow_cast( + temperature, + 'RunEndEncoded("run_ends": non-null Int32, "values": Int64)' + ) + RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING + ) AS windowed_sum +FROM sensor_readings +ORDER BY temperature; diff --git a/datafusion/sqllogictest/test_files/simplify_expr.slt b/datafusion/sqllogictest/test_files/simplify_expr.slt index 158096328e960..b24a1a6cff975 100644 --- a/datafusion/sqllogictest/test_files/simplify_expr.slt +++ b/datafusion/sqllogictest/test_files/simplify_expr.slt @@ -165,6 +165,60 @@ physical_plan 01)ProjectionExec: expr=[column1@0 = 1 as opt1, column1@0 = 2 AND column1@0 != 2 as noopt1, column1@0 = 4 as opt2, column1@0 != 5 AND column1@0 = 5 as noopt2] 02)--DataSourceExec: partitions=1, partition_sizes=[1] +# Dedup commutative AND/OR operands in join conditions +# (canonicalization is disabled for joins, so `A = B` vs `B = A` must be +# recognized directly). See https://github.com/apache/datafusion/issues/14943 + +statement ok +create table t1(a int) as values (1), (2); + +statement ok +create table t2(b int) as values (1), (3); + +# `t1.a = t2.b OR t2.b = t1.a` simplifies to `t1.a = t2.b`, allowing a +# HashJoin instead of a NestedLoopJoin +query TT +explain select * from t1 join t2 on t1.a = t2.b or t2.b = t1.a; +---- +logical_plan +01)Inner Join: t1.a = t2.b +02)--TableScan: t1 projection=[a] +03)--TableScan: t2 projection=[b] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, b@0)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +query II +select * from t1 join t2 on t1.a = t2.b or t2.b = t1.a; +---- +1 1 + +# swapped commutative operands (`t1.a + t2.b` vs `t2.b + t1.a`) dedup within +# a non-equi join filter +query TT +explain select * from t1 join t2 on t1.a > t2.b and (t1.a + t2.b > 1 or t2.b + t1.a > 1); +---- +logical_plan +01)Inner Join: Filter: t1.a > t2.b AND t1.a + t2.b > Int32(1) +02)--TableScan: t1 projection=[a] +03)--TableScan: t2 projection=[b] +physical_plan +01)NestedLoopJoinExec: join_type=Inner, filter=a@0 > b@1 AND a@0 + b@1 > 1 +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +query II +select * from t1 join t2 on t1.a > t2.b and (t1.a + t2.b > 1 or t2.b + t1.a > 1); +---- +2 1 + +statement ok +drop table t1; + +statement ok +drop table t2; + # Identity Date cast in a comparison predicate. # `cast(d AS date)` where `d` is already Date32 is an identity cast and should # fold away, so the predicate compares against the bare column `d`. This enables diff --git a/datafusion/sqllogictest/test_files/simplify_predicates.slt b/datafusion/sqllogictest/test_files/simplify_predicates.slt index 44fdedc9c8e1d..c36b0c864d592 100644 --- a/datafusion/sqllogictest/test_files/simplify_predicates.slt +++ b/datafusion/sqllogictest/test_files/simplify_predicates.slt @@ -242,5 +242,31 @@ logical_plan 02)--TableScan: test_data projection=[int_col, float_col, str_col, date_col, bool_col] +# x = 'apple' AND 'apple' = x is one condition, not a contradiction +query TT +EXPLAIN SELECT * FROM test_data WHERE str_col = 'apple' AND 'apple' = str_col; +---- +logical_plan +01)Filter: test_data.str_col = Utf8View("apple") +02)--TableScan: test_data projection=[int_col, float_col, str_col, date_col, bool_col] + +# x = 'apple' AND 'banana' = x is still impossible +query TT +EXPLAIN SELECT * FROM test_data WHERE str_col = 'apple' AND 'banana' = str_col; +---- +logical_plan EmptyRelation: rows=0 + statement ok set datafusion.explain.logical_plan_only=false; + +statement ok +CREATE TABLE fruit (name VARCHAR) AS VALUES ('apple'), ('banana'); + +query T +SELECT name FROM fruit WHERE name = 'apple' AND 'apple' = name; +---- +apple + +query T +SELECT name FROM fruit WHERE name = 'apple' AND 'banana' = name; +---- diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_batches.slt b/datafusion/sqllogictest/test_files/sort_merge_join_batches.slt new file mode 100644 index 0000000000000..eb26d0a715104 --- /dev/null +++ b/datafusion/sqllogictest/test_files/sort_merge_join_batches.slt @@ -0,0 +1,65 @@ +# 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. + +# Plan proof for sort_merge_join_matrix.slt's cross-batch coverage. SMJ sorts +# both inputs and SortExec re-batches to batch_size, so equal-key runs span +# batches even from a VALUES source. Not a matrix file: it fixes batch_size=1 to +# assert (via EXPLAIN ANALYZE) that the join then receives one row per batch +# (input_batches=16 = input_rows), i.e. the runs of five 1s and three 2s are +# split across batches. Fixtures match sort_merge_join_matrix.slt. + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +CREATE TABLE smj_dup_l(k INT) AS VALUES + (1), (1), (1), (1), (1), (2), (2), (2), (9); + +statement ok +CREATE TABLE smj_dup_r(k INT) AS VALUES + (1), (1), (1), (1), (2), (2), (3); + +query TT +EXPLAIN ANALYZE SELECT count(*) FROM smj_dup_l l JOIN smj_dup_r r ON l.k = r.k; +---- +Plan with Metrics +01) +02) +03) +04) +05)SortMergeJoinExec:input_batches=16, input_rows=16 + + +query I +SELECT count(*) FROM smj_dup_l l JOIN smj_dup_r r ON l.k = r.k; +---- +26 + +statement ok +RESET datafusion.execution.batch_size; + +statement ok +RESET datafusion.optimizer.prefer_hash_join; diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt b/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt new file mode 100644 index 0000000000000..7f1857782cb38 --- /dev/null +++ b/datafusion/sqllogictest/test_files/sort_merge_join_matrix.slt @@ -0,0 +1,315 @@ +# 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. + +# Sort-merge-join correctness across a config matrix. Every equijoin runs on +# SortMergeJoin (prefer_hash_join=false) and HashJoin (true), so sweeping +# {true,false} x batch_size {1,2,100,8192} requires them to agree row-for-row. +# batch_size=1 forces equal-key runs to span single-row batches (historically +# bug-prone). SMJ needs target_partitions>1 and repartition_joins, set below (not +# swept). Null-aware NOT IN always uses HashJoin, so it is out of scope (see +# null_aware_anti_join.slt). Matrix rules: no EXPLAIN, no in-file SET of a swept +# knob, rowsort every multi-row query. + +# configMatrix: datafusion.optimizer.prefer_hash_join=true,false +# configMatrix: datafusion.execution.batch_size=1,2,100,8192 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +# ------------------------------------------------------------------ +# Fixtures: duplicate keys, an unmatched key on each side, and a NULL key on +# each side (a NULL key never matches under default equality). +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_l(k INT, v VARCHAR) AS VALUES + (1, 'a'), (1, 'b'), (2, 'c'), (3, 'd'), (NULL, 'e'); + +statement ok +CREATE TABLE smj_r(k INT, w VARCHAR) AS VALUES + (1, 'x'), (1, 'y'), (2, 'z'), (4, 'q'), (NULL, 'n'); + +# ------------------------------------------------------------------ +# Basic: Inner / Left / Right / Full on a single equi key +# ------------------------------------------------------------------ +# INNER: k=1 gives 2x2, k=2 gives 1x1; k=3, k=4 and both NULLs are unmatched. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l JOIN smj_r r ON l.k = r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z + +# LEFT: the INNER rows plus unmatched left rows (k=3, NULL) padded with NULL. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l LEFT JOIN smj_r r ON l.k = r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z +3 d NULL +NULL e NULL + +# RIGHT: the INNER rows plus unmatched right rows (k=4, NULL) padded with NULL. +query TIT rowsort +SELECT l.v, r.k, r.w FROM smj_l l RIGHT JOIN smj_r r ON l.k = r.k; +---- +NULL 4 q +NULL NULL n +a 1 x +a 1 y +b 1 x +b 1 y +c 2 z + +# FULL: INNER rows plus every unmatched row from both sides. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l FULL JOIN smj_r r ON l.k = r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z +3 d NULL +NULL NULL n +NULL NULL q +NULL e NULL + +# ------------------------------------------------------------------ +# Existence joins over an equi key: LeftSemi / LeftAnti / RightSemi / RightAnti +# ------------------------------------------------------------------ +# LeftSemi (EXISTS): left rows whose key matches, each kept once (k=1 twice +# because the left side has two such rows, not because the right side does). +query IT rowsort +SELECT l.k, l.v FROM smj_l l WHERE EXISTS (SELECT 1 FROM smj_r r WHERE l.k = r.k); +---- +1 a +1 b +2 c + +# LeftAnti (NOT EXISTS): the complementary left rows, including the NULL key. +query IT rowsort +SELECT l.k, l.v FROM smj_l l WHERE NOT EXISTS (SELECT 1 FROM smj_r r WHERE l.k = r.k); +---- +3 d +NULL e + +# RightSemi: right rows whose key matches some left row. +query IT rowsort +SELECT r.k, r.w FROM smj_l l RIGHT SEMI JOIN smj_r r ON l.k = r.k; +---- +1 x +1 y +2 z + +# RightAnti: the complementary right rows, including the NULL key. +query IT rowsort +SELECT r.k, r.w FROM smj_l l RIGHT ANTI JOIN smj_r r ON l.k = r.k; +---- +4 q +NULL n + +# ------------------------------------------------------------------ +# Null-safe equality: `IS NOT DISTINCT FROM` makes NULL match NULL +# ------------------------------------------------------------------ +# The INNER result gains the (NULL,e)-(NULL,n) pair that plain `=` drops. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l JOIN smj_r r ON l.k IS NOT DISTINCT FROM r.k; +---- +1 a x +1 a y +1 b x +1 b y +2 c z +NULL e n + +# ------------------------------------------------------------------ +# Multiple equi keys +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_mk_l(k1 INT, k2 INT, v INT) AS VALUES + (1, 1, 10), (1, 2, 20), (2, 2, 30), (NULL, 1, 40); + +statement ok +CREATE TABLE smj_mk_r(k1 INT, k2 INT, w INT) AS VALUES + (1, 1, 100), (1, 2, 200), (2, 2, 300), (1, 3, 400); + +query IIII rowsort +SELECT l.k1, l.k2, l.v, r.w +FROM smj_mk_l l JOIN smj_mk_r r ON l.k1 = r.k1 AND l.k2 = r.k2; +---- +1 1 10 100 +1 2 20 200 +2 2 30 300 + +query IIII rowsort +SELECT l.k1, l.k2, l.v, r.w +FROM smj_mk_l l LEFT JOIN smj_mk_r r ON l.k1 = r.k1 AND l.k2 = r.k2; +---- +1 1 10 100 +1 2 20 200 +2 2 30 300 +NULL 1 40 NULL + +# ------------------------------------------------------------------ +# Equi key plus a non-equi join filter in the ON clause +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_a(k INT, x INT) AS VALUES (1, 10), (1, 20), (2, 30), (3, 40); + +statement ok +CREATE TABLE smj_b(k INT, y INT) AS VALUES (1, 15), (1, 25), (2, 5), (4, 50); + +# INNER: only k=1 has matches, kept where x < y. +query IIII rowsort +SELECT l.k, l.x, r.k, r.y FROM smj_a l JOIN smj_b r ON l.k = r.k AND l.x < r.y; +---- +1 10 1 15 +1 10 1 25 +1 20 1 25 + +# LEFT: rows whose key never matches, or matches but fails the filter, keep the +# left row with NULLs (k=2's x=30 fails 30<5; k=3 has no right key). +query III rowsort +SELECT l.k, l.x, r.y FROM smj_a l LEFT JOIN smj_b r ON l.k = r.k AND l.x < r.y; +---- +1 10 15 +1 10 25 +1 20 25 +2 30 NULL +3 40 NULL + +# FULL with the same filter: unmatched rows from both sides survive. +query IIII rowsort +SELECT l.k, l.x, r.k, r.y FROM smj_a l FULL JOIN smj_b r ON l.k = r.k AND l.x < r.y; +---- +1 10 1 15 +1 10 1 25 +1 20 1 25 +2 30 NULL NULL +3 40 NULL NULL +NULL NULL 2 5 +NULL NULL 4 50 + +# ------------------------------------------------------------------ +# Non-integer (string) key +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_sk_l(k VARCHAR, v INT) AS VALUES ('a', 1), ('b', 2), ('a', 3); + +statement ok +CREATE TABLE smj_sk_r(k VARCHAR, w INT) AS VALUES ('a', 10), ('c', 20); + +query TII rowsort +SELECT l.k, l.v, r.w FROM smj_sk_l l JOIN smj_sk_r r ON l.k = r.k; +---- +a 1 10 +a 3 10 + +query TII rowsort +SELECT l.k, l.v, r.w FROM smj_sk_l l FULL JOIN smj_sk_r r ON l.k = r.k; +---- +NULL NULL 20 +a 1 10 +a 3 10 +b 2 NULL + +# ------------------------------------------------------------------ +# Empty inputs +# ------------------------------------------------------------------ +statement ok +CREATE TABLE smj_empty(k INT, w VARCHAR); + +# INNER against an empty side is empty; LEFT keeps every left row with NULLs. +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l JOIN smj_empty r ON l.k = r.k; +---- + +query ITT rowsort +SELECT l.k, l.v, r.w FROM smj_l l LEFT JOIN smj_empty r ON l.k = r.k; +---- +1 a NULL +1 b NULL +2 c NULL +3 d NULL +NULL e NULL + +# Empty left side with a RIGHT join keeps every right row with NULLs. +query IIT rowsort +SELECT l.k, r.k, r.w FROM smj_empty l RIGHT JOIN smj_r r ON l.k = r.k; +---- +NULL 1 x +NULL 1 y +NULL 2 z +NULL 4 q +NULL NULL n + +# ------------------------------------------------------------------ +# Duplicate keys across batch boundaries (verified by count) +# ------------------------------------------------------------------ +# SortExec re-batches to batch_size, so equal-key runs span batches (proof in +# sort_merge_join_batches.slt). Left: five k=1, three k=2, one unmatched k=9; +# right: four k=1, two k=2, one unmatched k=3. INNER = 5*4 + 3*2 = 26. +statement ok +CREATE TABLE smj_dup_l(k INT) AS VALUES + (1), (1), (1), (1), (1), (2), (2), (2), (9); + +statement ok +CREATE TABLE smj_dup_r(k INT) AS VALUES + (1), (1), (1), (1), (2), (2), (3); + +query I +SELECT count(*) FROM smj_dup_l l JOIN smj_dup_r r ON l.k = r.k; +---- +26 + +# LEFT = INNER + unmatched left k=9 = 27. +query I +SELECT count(*) FROM smj_dup_l l LEFT JOIN smj_dup_r r ON l.k = r.k; +---- +27 + +# RIGHT = INNER + unmatched right k=3 = 27. +query I +SELECT count(*) FROM smj_dup_l l RIGHT JOIN smj_dup_r r ON l.k = r.k; +---- +27 + +# FULL = INNER + both unmatched = 28. +query I +SELECT count(*) FROM smj_dup_l l FULL JOIN smj_dup_r r ON l.k = r.k; +---- +28 + +# LeftSemi = left rows with a match = 5 + 3 = 8. +query I +SELECT count(*) FROM smj_dup_l l WHERE EXISTS (SELECT 1 FROM smj_dup_r r WHERE l.k = r.k); +---- +8 + +# LeftAnti = unmatched left rows = 1 (k=9). +query I +SELECT count(*) FROM smj_dup_l l WHERE NOT EXISTS (SELECT 1 FROM smj_dup_r r WHERE l.k = r.k); +---- +1 diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index 8c0114b90ee13..982a87d6d5e83 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -17,8 +17,11 @@ use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer}; use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts}; -use datafusion::common::{DFSchemaRef, not_impl_err}; -use datafusion::logical_expr::{Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder}; +use datafusion::common::{Column, DFSchemaRef, internal_err, not_impl_err}; +use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::{ + Aggregate, Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder, +}; use substrait::proto::AggregateRel; use substrait::proto::aggregate_function::AggregationInvocation; use substrait::proto::aggregate_rel::Grouping; @@ -122,12 +125,51 @@ pub async fn from_aggregate_rel( .map(|e| name_tracker.get_uniquely_named_expr(e)) .collect::, _>>()?; - input.aggregate(group_exprs, aggr_exprs)?.build() + let plan = input.aggregate(group_exprs, aggr_exprs)?.build()?; + if agg.groupings.len() > 1 { + reorder_grouping_set_output(plan, agg.measures.len()) + } else { + Ok(plan) + } } else { not_impl_err!("Aggregate without an input is not valid") } } +/// Reorders DataFusion's `[groups, grouping_id, measures]` aggregate schema to +/// Substrait's direct output order of `[groups, measures, grouping_id]`. +fn reorder_grouping_set_output( + plan: LogicalPlan, + measure_count: usize, +) -> datafusion::common::Result { + let exprs: Vec = { + let schema = plan.schema(); + let Some(grouping_id_index) = + schema.index_of_column_by_name(None, Aggregate::INTERNAL_GROUPING_ID) + else { + return internal_err!( + "Grouping set aggregate schema is missing {}", + Aggregate::INTERNAL_GROUPING_ID + ); + }; + if grouping_id_index + measure_count + 1 != schema.fields().len() { + return internal_err!( + "Grouping set aggregate schema has {} fields after {}, expected {} measures", + schema.fields().len() - grouping_id_index - 1, + Aggregate::INTERNAL_GROUPING_ID, + measure_count + ); + } + + (0..grouping_id_index) + .chain(grouping_id_index + 1..schema.fields().len()) + .chain(std::iter::once(grouping_id_index)) + .map(|index| Expr::Column(Column::from(schema.qualified_field(index)))) + .collect() + }; + project(plan, exprs) +} + #[expect(deprecated)] async fn from_substrait_grouping( consumer: &impl SubstraitConsumer, diff --git a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs index dec94b0422257..7b6c113ccec0f 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/aggregate_rel.rs @@ -24,7 +24,8 @@ use datafusion::logical_expr::utils::powerset; use datafusion::logical_expr::{Aggregate, Distinct, Expr, GroupingSet}; use substrait::proto::aggregate_rel::{Grouping, Measure}; use substrait::proto::rel::RelType; -use substrait::proto::{AggregateRel, Expression, Rel}; +use substrait::proto::rel_common::EmitKind; +use substrait::proto::{AggregateRel, Expression, Rel, RelCommon, rel_common}; pub fn from_aggregate( producer: &mut impl SubstraitProducer, @@ -38,10 +39,12 @@ pub fn from_aggregate( .iter() .map(|e| to_substrait_agg_measure(producer, e, agg.input.schema())) .collect::>>()?; + let common = (groupings.len() > 1) + .then(|| grouping_set_output_mapping(grouping_expressions.len(), measures.len())); Ok(Box::new(Rel { rel_type: Some(RelType::Aggregate(Box::new(AggregateRel { - common: None, + common, input: Some(input), grouping_expressions, groupings, @@ -51,6 +54,22 @@ pub fn from_aggregate( })) } +/// Maps Substrait's `[groups, measures, grouping_id]` direct output to +/// DataFusion's `[groups, grouping_id, measures]` aggregate schema. +fn grouping_set_output_mapping(grouping_count: usize, measure_count: usize) -> RelCommon { + let grouping_id_index = grouping_count + measure_count; + let output_mapping = (0..grouping_count) + .chain(std::iter::once(grouping_id_index)) + .chain(grouping_count..grouping_id_index) + .map(|index| index as i32) + .collect(); + RelCommon { + emit_kind: Some(EmitKind::Emit(rel_common::Emit { output_mapping })), + hint: None, + advanced_extension: None, + } +} + pub fn from_distinct( producer: &mut impl SubstraitProducer, distinct: &Distinct, @@ -165,8 +184,12 @@ pub fn parse_flat_grouping_exprs( for e in exprs { let rex = producer.handle_expr(e, schema)?; grouping_expressions.push(rex.clone()); - ref_group_exprs.push(rex); - expression_references.push((ref_group_exprs.len() - 1) as u32); + let reference = ref_group_exprs.iter().position(|existing| existing == &rex); + let reference = reference.unwrap_or_else(|| { + ref_group_exprs.push(rex); + ref_group_exprs.len() - 1 + }); + expression_references.push(reference as u32); } #[expect(deprecated)] Ok(Grouping { diff --git a/datafusion/substrait/tests/cases/aggregation_tests.rs b/datafusion/substrait/tests/cases/aggregation_tests.rs index 92a41850b208d..e572023f17a92 100644 --- a/datafusion/substrait/tests/cases/aggregation_tests.rs +++ b/datafusion/substrait/tests/cases/aggregation_tests.rs @@ -68,4 +68,42 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn multiple_grouping_sets_follow_substrait_output_order() -> Result<()> { + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json", + ); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + assert_snapshot!( + plan, + @r" + Projection: c0, c1, sum(c0) AS summation + Aggregate: groupBy=[[GROUPING SETS ((c0), (c1), (c0, c1))]], aggr=[[sum(c0)]] + Values: (Int64(1), Int64(10)), (Int64(1), Int64(20)), (Int64(2), Int64(10)) + " + ); + + let results = DataFrame::new(ctx.state(), plan).collect().await?; + datafusion::assert_batches_sorted_eq!( + [ + "+----+----+-----------+", + "| c0 | c1 | summation |", + "+----+----+-----------+", + "| | 10 | 3 |", + "| | 20 | 1 |", + "| 1 | | 2 |", + "| 1 | 10 | 1 |", + "| 1 | 20 | 1 |", + "| 2 | | 2 |", + "| 2 | 10 | 2 |", + "+----+----+-----------+", + ], + &results + ); + + Ok(()) + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index c9f874dd9b095..3716b0feba3cc 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -344,10 +344,59 @@ async fn aggregate_multiple_keys() -> Result<()> { #[tokio::test] async fn aggregate_grouping_sets() -> Result<()> { - roundtrip( - "SELECT a, c, d, avg(b) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())", + let ctx = create_context().await?; + let proto = roundtrip_with_ctx( + "SELECT a, c, d, avg(b), sum(e) FROM data GROUP BY GROUPING SETS ((a, c), (a), (d), ())", + ctx.clone(), ) - .await + .await?; + + let Some(plan_rel::RelType::Root(root)) = &proto.relations[0].rel_type else { + panic!("expected root relation"); + }; + let Some(RelType::Project(project)) = &root.input.as_ref().unwrap().rel_type else { + panic!("expected project relation"); + }; + let Some(RelType::Aggregate(aggregate)) = &project.input.as_ref().unwrap().rel_type + else { + panic!("expected aggregate relation"); + }; + + assert_eq!(aggregate.grouping_expressions.len(), 3); + assert_eq!(aggregate.groupings[0].expression_references, [0, 1]); + assert_eq!(aggregate.groupings[1].expression_references, [0]); + assert_eq!(aggregate.groupings[2].expression_references, [2]); + assert!(aggregate.groupings[3].expression_references.is_empty()); + let output_mapping = match aggregate + .common + .as_ref() + .and_then(|common| common.emit_kind.as_ref()) + { + Some(substrait::proto::rel_common::EmitKind::Emit(emit)) => &emit.output_mapping, + _ => panic!("expected aggregate output mapping"), + }; + assert_eq!(output_mapping, &[0, 1, 2, 5, 3, 4]); + + let plan = from_substrait_plan(&ctx.state(), &proto).await?; + let results = DataFrame::new(ctx.state(), plan).collect().await?; + datafusion::assert_batches_sorted_eq!( + [ + "+---+------------+-------+-------------+-------------+", + "| a | c | d | avg(data.b) | sum(data.e) |", + "+---+------------+-------+-------------+-------------+", + "| | | | 3.250000 | 6442450943 |", + "| | | false | 2.000000 | 4294967295 |", + "| | | true | 4.500000 | 2147483648 |", + "| 1 | | | 2.000000 | 4294967295 |", + "| 1 | 2020-01-01 | | 2.000000 | 4294967295 |", + "| 3 | | | 4.500000 | 2147483648 |", + "| 3 | 2020-01-01 | | 4.500000 | 2147483648 |", + "+---+------------+-------+-------------+-------------+", + ], + &results + ); + + Ok(()) } #[tokio::test] diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json new file mode 100644 index 0000000000000..3779ee32c6ab8 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_groupings/multiple_groupings.json @@ -0,0 +1,127 @@ +{ + "extensionUris": [ + { + "extensionUriAnchor": 1, + "uri": "https://github.com/substrait-io/substrait/blob/main/extensions/functions_arithmetic.yaml" + } + ], + "extensions": [ + { + "extensionFunction": { + "extensionUriReference": 1, + "functionAnchor": 1, + "name": "sum:i64" + } + } + ], + "relations": [ + { + "root": { + "input": { + "aggregate": { + "common": { + "emit": { + "outputMapping": [0, 1, 2] + } + }, + "input": { + "read": { + "baseSchema": { + "names": ["c0", "c1"], + "struct": { + "nullability": "NULLABILITY_REQUIRED", + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + }, + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ] + } + }, + "common": { + "direct": {} + }, + "virtualTable": { + "expressions": [ + { + "fields": [ + { "literal": { "i64": "1", "nullable": false } }, + { "literal": { "i64": "10", "nullable": false } } + ] + }, + { + "fields": [ + { "literal": { "i64": "1", "nullable": false } }, + { "literal": { "i64": "20", "nullable": false } } + ] + }, + { + "fields": [ + { "literal": { "i64": "2", "nullable": false } }, + { "literal": { "i64": "10", "nullable": false } } + ] + } + ] + } + } + }, + "groupingExpressions": [ + { + "selection": { + "directReference": { "structField": {} }, + "rootReference": {} + } + }, + { + "selection": { + "directReference": { "structField": { "field": 1 } }, + "rootReference": {} + } + } + ], + "groupings": [ + { "expressionReferences": [0] }, + { "expressionReferences": [1] }, + { "expressionReferences": [0, 1] } + ], + "measures": [ + { + "measure": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { "structField": {} }, + "rootReference": {} + } + } + } + ], + "functionReference": 1, + "invocation": "AGGREGATION_INVOCATION_ALL", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT" + } + } + ] + } + }, + "names": ["c0", "c1", "summation"] + } + } + ], + "version": { + "minorNumber": 29, + "producer": "substrait-go v4.2.0" + } +} diff --git a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json index a43843bd0fe9f..09422c7001a55 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json @@ -13,8 +13,8 @@ }, "devDependencies": { "copy-webpack-plugin": "14.0.0", - "webpack": "5.109.2", - "webpack-cli": "7.2.2", + "webpack": "5.110.1", + "webpack-cli": "7.2.3", "webpack-dev-server": "6.0.0" } }, @@ -36,6 +36,7 @@ "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz", "integrity": "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==", "dev": true, + "license": "MIT", "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.0", "@jridgewell/trace-mapping": "^0.3.24" @@ -46,6 +47,7 @@ "resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz", "integrity": "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==", "dev": true, + "license": "MIT", "engines": { "node": ">=6.0.0" } @@ -55,22 +57,25 @@ "resolved": "https://registry.npmjs.org/@jridgewell/source-map/-/source-map-0.3.11.tgz", "integrity": "sha512-ZMp1V8ZFcPG5dIWnQLr3NSI1MiCU7UETdS/A0G8V/XWHvJv3ZsFqutJn1Y5RPmAPX6F3BiE397OqveU/9NCuIA==", "dev": true, + "license": "MIT", "dependencies": { "@jridgewell/gen-mapping": "^0.3.5", "@jridgewell/trace-mapping": "^0.3.25" } }, "node_modules/@jridgewell/sourcemap-codec": { - "version": "1.5.5", - "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", - "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", - "dev": true + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz", + "integrity": "sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==", + "dev": true, + "license": "MIT" }, "node_modules/@jridgewell/trace-mapping": { "version": "0.3.31", "resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz", "integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==", "dev": true, + "license": "MIT", "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" @@ -670,6 +675,7 @@ "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.18.0.tgz", "integrity": "sha512-lGq+9yr1/GuAWaVYIHRjvvySG5/4VfKIvC8EWxStPdcDh/Ka7FG3twP6v4d5BkravUilhIAsG4Qj83t02LWUPQ==", "dev": true, + "license": "MIT", "bin": { "acorn": "bin/acorn" }, @@ -755,12 +761,16 @@ "dev": true }, "node_modules/baseline-browser-mapping": { - "version": "2.9.19", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.9.19.tgz", - "integrity": "sha512-ipDqC8FrAl/76p2SSWKSI+H9tFwm7vYqXQrItCuiVPt26Km0jS+NzSsBWAaBusvSbQcfJG+JitdMm+wZAgTYqg==", + "version": "2.11.20", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.20.tgz", + "integrity": "sha512-H0ulySigv6icDJ1F7SjtdCD6PrhTpdYCmP0CactWy1+ekh0AFd0o1Wn5T8b+hnTmdBx19u9yhL6wvCylXMY7zw==", "dev": true, + "license": "Apache-2.0", "bin": { - "baseline-browser-mapping": "dist/cli.js" + "baseline-browser-mapping": "dist/cli.cjs" + }, + "engines": { + "node": ">=6.0.0" } }, "node_modules/batch": { @@ -830,9 +840,9 @@ } }, "node_modules/browserslist": { - "version": "4.28.1", - "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.1.tgz", - "integrity": "sha512-ZC5Bd0LgJXgwGqUknZY/vkUQ04r8NXnJZ3yYi4vDmSiZmC/pdSN0NbNRPxZpbtO4uAfDUAFffO8IZoM3Gj8IkA==", + "version": "4.28.8", + "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.8.tgz", + "integrity": "sha512-V2NpofLblG64mfOtSgDhOJESZEGogzDMBv/q+W6oc4LXWP/q75eOXoOaaOu1EOadB9U4Bwx/e0yzbvwKH8zalA==", "dev": true, "funding": [ { @@ -848,12 +858,13 @@ "url": "https://github.com/sponsors/ai" } ], + "license": "MIT", "dependencies": { - "baseline-browser-mapping": "^2.9.0", - "caniuse-lite": "^1.0.30001759", - "electron-to-chromium": "^1.5.263", - "node-releases": "^2.0.27", - "update-browserslist-db": "^1.2.0" + "baseline-browser-mapping": "^2.11.12", + "caniuse-lite": "^1.0.30001809", + "electron-to-chromium": "^1.5.402", + "node-releases": "^2.0.53", + "update-browserslist-db": "^1.3.0" }, "bin": { "browserslist": "cli.js" @@ -866,7 +877,8 @@ "version": "1.1.2", "resolved": "https://registry.npmjs.org/buffer-from/-/buffer-from-1.1.2.tgz", "integrity": "sha512-E+XQCRwSbaaiChtv6k6Dwgc+bx+Bs6vuKJHHl5kox/BaKbhiXzqQOwK4cO22yElGp2OCmjwVhT3HmxgyPGnJfQ==", - "dev": true + "dev": true, + "license": "MIT" }, "node_modules/bundle-name": { "version": "4.1.0", @@ -931,9 +943,9 @@ } }, "node_modules/caniuse-lite": { - "version": "1.0.30001768", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001768.tgz", - "integrity": "sha512-qY3aDRZC5nWPgHUgIB84WL+nySuo19wk0VJpp/XI9T34lrvkyhRvNVOFJOp2kxClQhiFBu+TaUSudf6oa3vkSA==", + "version": "1.0.30001810", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz", + "integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==", "dev": true, "funding": [ { @@ -948,7 +960,8 @@ "type": "github", "url": "https://github.com/sponsors/ai" } - ] + ], + "license": "CC-BY-4.0" }, "node_modules/chokidar": { "version": "5.0.0", @@ -995,7 +1008,8 @@ "version": "2.20.3", "resolved": "https://registry.npmjs.org/commander/-/commander-2.20.3.tgz", "integrity": "sha512-GpVkmM8vF2vQUkj2LvZmD35JxeJOLCwJ9cUkugyk2nuhbv3+mJvpLYYt+0+USMxE+oj+ey/lJEnhZw75x/OMcQ==", - "dev": true + "dev": true, + "license": "MIT" }, "node_modules/compressible": { "version": "2.0.18", @@ -1283,10 +1297,11 @@ "dev": true }, "node_modules/electron-to-chromium": { - "version": "1.5.286", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.286.tgz", - "integrity": "sha512-9tfDXhJ4RKFNerfjdCcZfufu49vg620741MNs26a9+bhLThdB+plgMeou98CAaHu/WATj2iHOOHTp1hWtABj2A==", - "dev": true + "version": "1.5.419", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.419.tgz", + "integrity": "sha512-nHMPn8x4yCxCI0iSnL+LlHL5sUoUfjLXkcRIagZ4GBdrfFLFaiLNvzJWbJqZhFT9IAhw5tUSNlhggWN+otvp/A==", + "dev": true, + "license": "ISC" }, "node_modules/encodeurl": { "version": "2.0.0", @@ -1363,6 +1378,7 @@ "resolved": "https://registry.npmjs.org/escalade/-/escalade-3.2.0.tgz", "integrity": "sha512-WUj2qlxaQtO4g6Pq5c29GTcWGDyd8itL8zTlipgECz3JesAiiOKotd8JU6otB3PACgG6xkJUyVhboMS+bje/jA==", "dev": true, + "license": "MIT", "engines": { "node": ">=6" } @@ -1373,49 +1389,6 @@ "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", "dev": true }, - "node_modules/eslint-scope": { - "version": "5.1.1", - "resolved": "https://registry.npmjs.org/eslint-scope/-/eslint-scope-5.1.1.tgz", - "integrity": "sha512-2NxwbF/hZ0KpepYN0cNbo+FN6XoK7GaHlQhgx/hIZl6Va0bF45RQOOwhLIy8lQDbuCiadSLCBnH2CFYquit5bw==", - "dev": true, - "dependencies": { - "esrecurse": "^4.3.0", - "estraverse": "^4.1.1" - }, - "engines": { - "node": ">=8.0.0" - } - }, - "node_modules/esrecurse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/esrecurse/-/esrecurse-4.3.0.tgz", - "integrity": "sha512-KmfKL3b6G+RXvP8N1vr3Tq1kL/oCFgn2NYXEtqP8/L3pKapUA4G8cFVaoF3SU323CD4XypR/ffioHmkti6/Tag==", - "dev": true, - "dependencies": { - "estraverse": "^5.2.0" - }, - "engines": { - "node": ">=4.0" - } - }, - "node_modules/esrecurse/node_modules/estraverse": { - "version": "5.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-5.3.0.tgz", - "integrity": "sha512-MMdARuVEQziNTeJD8DgMqmhwR11BRQ/cBP+pLtYdSTnf3MIO8fFeiINEbX36ZdNlfU/7A9f3gUw49B3oQsvwBA==", - "dev": true, - "engines": { - "node": ">=4.0" - } - }, - "node_modules/estraverse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-4.3.0.tgz", - "integrity": "sha512-39nnKffWz8xN1BU/2c79n9nB9HDzo0niYUqx6xyqUnyoAnQyyWpOTdZEeiCch8BBu515t4wp9ZmgVfVhn9EBpw==", - "dev": true, - "engines": { - "node": ">=4.0" - } - }, "node_modules/etag": { "version": "1.8.1", "resolved": "https://registry.npmjs.org/etag/-/etag-1.8.1.tgz", @@ -1509,9 +1482,9 @@ "dev": true }, "node_modules/fast-uri": { - "version": "3.1.5", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", - "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", + "version": "3.1.7", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.7.tgz", + "integrity": "sha512-dOvZVzjdZdz7phd9v6jCbwxrBW3fK6n8Rc0CtdmM4bumzMnxywBYhuph6J819RRw/ku+rLbelwfMunktuzVVHg==", "dev": true, "funding": [ { @@ -1522,7 +1495,8 @@ "type": "opencollective", "url": "https://opencollective.com/fastify" } - ] + ], + "license": "BSD-3-Clause" }, "node_modules/fill-range": { "version": "7.1.1", @@ -1678,6 +1652,7 @@ "resolved": "https://registry.npmjs.org/has-flag/-/has-flag-4.0.0.tgz", "integrity": "sha512-EykJT/Q1KjTWctppgIAgfSO0tKVuZUjhgMr17kqTumMl6Afv3EISleU7qZUzoXDFTAHTDC4NOoG/ZxU3EvlMPQ==", "dev": true, + "license": "MIT", "engines": { "node": ">=8" } @@ -1980,6 +1955,7 @@ "resolved": "https://registry.npmjs.org/jest-worker/-/jest-worker-27.5.1.tgz", "integrity": "sha512-7vuh85V5cdDofPyxn58nrPjBktZo0u9x1g8WtjQol+jZDaE+fhN+cIvTj11GndBnMnyfrUOG1sZQxCdjKh+DKg==", "dev": true, + "license": "MIT", "dependencies": { "@types/node": "*", "merge-stream": "^2.0.0", @@ -2557,7 +2533,8 @@ "version": "2.0.0", "resolved": "https://registry.npmjs.org/merge-stream/-/merge-stream-2.0.0.tgz", "integrity": "sha512-abv/qOcuPfk3URPfDzmZU1LKmuw8kT+0nIHvKrKgFrwifol/doWcdA4ZqsWQ8ENrFKkd67Mfpo/LovbIUsbt3w==", - "dev": true + "dev": true, + "license": "MIT" }, "node_modules/micromatch": { "version": "4.0.8", @@ -2594,15 +2571,16 @@ } }, "node_modules/minimizer-webpack-plugin": { - "version": "5.6.1", - "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.6.1.tgz", - "integrity": "sha512-DoeAZz8Q1C1znwsUzej1fdoi4jCf7/+Em27ouLqfK/+3m8G+D7yDhUwrc3CNhjSzGUN1kn7Iv4sWmjflQHenpw==", + "version": "5.8.0", + "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.8.0.tgz", + "integrity": "sha512-2cT9+goJfBhtMz+gJqejSf09ClgmYhPhccRb/fb0ztVbixt0BkO8mRuI26FoPPuUNkRk6iEDryWAIouc5W8eJA==", "dev": true, + "license": "MIT", "dependencies": { - "@jridgewell/trace-mapping": "^0.3.25", + "@jridgewell/trace-mapping": "^0.3.31", "jest-worker": "^27.4.5", - "schema-utils": "^4.3.0", - "terser": "^5.31.1" + "schema-utils": "^4.3.3", + "terser": "^5.51.0" }, "engines": { "node": ">= 10.13.0" @@ -2689,10 +2667,14 @@ "dev": true }, "node_modules/node-releases": { - "version": "2.0.27", - "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.27.tgz", - "integrity": "sha512-nmh3lCkYZ3grZvqcCH+fjmQ7X+H0OeZgP40OierEaAptX4XofMh5kwNbWh7lBduUzCcV/8kZ+NDLCwm2iorIlA==", - "dev": true + "version": "2.0.54", + "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz", + "integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=18" + } }, "node_modules/normalize-path": { "version": "3.0.0", @@ -3449,6 +3431,7 @@ "resolved": "https://registry.npmjs.org/source-map/-/source-map-0.6.1.tgz", "integrity": "sha512-UjgapumWlbMhkBgzT7Ykc5YXUT46F0iKu8SGXq0bcwP5dz/h0Plj6enJqjz1Zbq2l5WaqYnrVbwWOWMyF3F47g==", "dev": true, + "license": "BSD-3-Clause", "engines": { "node": ">=0.10.0" } @@ -3458,6 +3441,7 @@ "resolved": "https://registry.npmjs.org/source-map-support/-/source-map-support-0.5.21.tgz", "integrity": "sha512-uBHU3L3czsIyYXKX88fdrGovxdSCoTGDRZ6SYXtSRxLZUzHg5P/66Ht6uoUlHu9EZod+inXhKo3qQgwXUT/y1w==", "dev": true, + "license": "MIT", "dependencies": { "buffer-from": "^1.0.0", "source-map": "^0.6.0" @@ -3477,6 +3461,7 @@ "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-8.1.1.tgz", "integrity": "sha512-MpUEN2OodtUzxvKQl72cUF7RQ5EiHsGvSsVG0ia9c5RbWGL2CI4C7EpPS8UTBIplnlzZiNuV56w+FuNxy3ty2Q==", "dev": true, + "license": "MIT", "dependencies": { "has-flag": "^4.0.0" }, @@ -3513,10 +3498,11 @@ } }, "node_modules/terser": { - "version": "5.50.0", - "resolved": "https://registry.npmjs.org/terser/-/terser-5.50.0.tgz", - "integrity": "sha512-CN9BVxWhgS/hRxtUMjtC2uRWSTcSfQFHMDWma6sKKfIivCD91sM+FOPfvwoaRMqCSrUpe1nv3jDamd9eEQ4y+w==", + "version": "5.51.2", + "resolved": "https://registry.npmjs.org/terser/-/terser-5.51.2.tgz", + "integrity": "sha512-bWnjSNscmuI+GJze6ZupnHP8G/cTcsJF+bXCeQknk2SHQsgbNJnLrqiH9jZ2W4STPVXH2mDKKRX3iwPhc9Cn/Q==", "dev": true, + "license": "BSD-2-Clause", "dependencies": { "@jridgewell/source-map": "^0.3.3", "acorn": "^8.15.0", @@ -3687,9 +3673,9 @@ } }, "node_modules/update-browserslist-db": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", - "integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz", + "integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==", "dev": true, "funding": [ { @@ -3705,6 +3691,7 @@ "url": "https://github.com/sponsors/ai" } ], + "license": "MIT", "dependencies": { "escalade": "^3.2.0", "picocolors": "^1.1.1" @@ -3738,10 +3725,11 @@ } }, "node_modules/webpack": { - "version": "5.109.2", - "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.109.2.tgz", - "integrity": "sha512-U9/cvLzxObKNEZ9+TtdqrHM5/9z3lgl2c+c4BzbqGxFQvQvBAq87yql5A8pQ+rrMbS496MZJeF5enVBndIy2hw==", + "version": "5.110.1", + "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.110.1.tgz", + "integrity": "sha512-gInQB+jxXxgnZyvPwuzT5NGQmECDqeu85oxcrjinrYHqPoBex0hCAN2SFTJVyPVrK0Pq9E44VFP+e89fAc10/w==", "dev": true, + "license": "MIT", "dependencies": { "@types/estree": "^1.0.8", "@types/json-schema": "^7.0.15", @@ -3753,11 +3741,10 @@ "chrome-trace-event": "^1.0.2", "enhanced-resolve": "^5.24.4", "es-module-lexer": "^2.1.0", - "eslint-scope": "5.1.1", "events": "^3.2.0", "graceful-fs": "^4.2.11", "mime-db": "^1.54.0", - "minimizer-webpack-plugin": "^5.6.1", + "minimizer-webpack-plugin": "^5.7.0", "neo-async": "^2.6.2", "schema-utils": "^4.3.3", "tapable": "^2.3.0", @@ -3781,10 +3768,11 @@ } }, "node_modules/webpack-cli": { - "version": "7.2.2", - "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.2.tgz", - "integrity": "sha512-lD0pALneslq8FfV+rwvm1BMW0AFAJrHHhNupAGN4asYjMvqrtRsenU4iKpiBo09gS4ntMxKGUxl9jhTEzVt0oA==", + "version": "7.2.3", + "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.3.tgz", + "integrity": "sha512-vDFU7jrfCctnN7jJQWPl+V26B51GLp11prVZXg50oeonsgeBzSTJEWmXkjHsOjTgMjlPOQM8GWLh33X9RN/0ow==", "dev": true, + "license": "MIT", "dependencies": { "@discoveryjs/json-ext": "^1.1.0", "commander": "^14.0.3", @@ -3808,7 +3796,7 @@ "peerDependencies": { "js-yaml": "^4.0.0 || ^5.0.0", "json5": "^2.2.3", - "toml": "^3.0.0 || ^4.0.0", + "toml": "^3.0.0 || ^4.0.0 || ^5.0.0", "webpack": "^5.101.0", "webpack-bundle-analyzer": "^4.0.0 || ^5.0.0", "webpack-dev-server": "^5.0.0 || ^6.0.0" @@ -4077,9 +4065,9 @@ } }, "@jridgewell/sourcemap-codec": { - "version": "1.5.5", - "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", - "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz", + "integrity": "sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==", "dev": true }, "@jridgewell/trace-mapping": { @@ -4742,9 +4730,9 @@ } }, "baseline-browser-mapping": { - "version": "2.9.19", - "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.9.19.tgz", - "integrity": "sha512-ipDqC8FrAl/76p2SSWKSI+H9tFwm7vYqXQrItCuiVPt26Km0jS+NzSsBWAaBusvSbQcfJG+JitdMm+wZAgTYqg==", + "version": "2.11.20", + "resolved": "https://registry.npmjs.org/baseline-browser-mapping/-/baseline-browser-mapping-2.11.20.tgz", + "integrity": "sha512-H0ulySigv6icDJ1F7SjtdCD6PrhTpdYCmP0CactWy1+ekh0AFd0o1Wn5T8b+hnTmdBx19u9yhL6wvCylXMY7zw==", "dev": true }, "batch": { @@ -4798,16 +4786,16 @@ } }, "browserslist": { - "version": "4.28.1", - "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.1.tgz", - "integrity": "sha512-ZC5Bd0LgJXgwGqUknZY/vkUQ04r8NXnJZ3yYi4vDmSiZmC/pdSN0NbNRPxZpbtO4uAfDUAFffO8IZoM3Gj8IkA==", + "version": "4.28.8", + "resolved": "https://registry.npmjs.org/browserslist/-/browserslist-4.28.8.tgz", + "integrity": "sha512-V2NpofLblG64mfOtSgDhOJESZEGogzDMBv/q+W6oc4LXWP/q75eOXoOaaOu1EOadB9U4Bwx/e0yzbvwKH8zalA==", "dev": true, "requires": { - "baseline-browser-mapping": "^2.9.0", - "caniuse-lite": "^1.0.30001759", - "electron-to-chromium": "^1.5.263", - "node-releases": "^2.0.27", - "update-browserslist-db": "^1.2.0" + "baseline-browser-mapping": "^2.11.12", + "caniuse-lite": "^1.0.30001809", + "electron-to-chromium": "^1.5.402", + "node-releases": "^2.0.53", + "update-browserslist-db": "^1.3.0" } }, "buffer-from": { @@ -4858,9 +4846,9 @@ } }, "caniuse-lite": { - "version": "1.0.30001768", - "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001768.tgz", - "integrity": "sha512-qY3aDRZC5nWPgHUgIB84WL+nySuo19wk0VJpp/XI9T34lrvkyhRvNVOFJOp2kxClQhiFBu+TaUSudf6oa3vkSA==", + "version": "1.0.30001810", + "resolved": "https://registry.npmjs.org/caniuse-lite/-/caniuse-lite-1.0.30001810.tgz", + "integrity": "sha512-TITQPUkaz+aVk5GL6NhOdwk1aEaNTSDPsGFWrTuhKGtjTF70jL/Oht2W4c6rXUe5fu7Ie19VIahAXHIIiWWNeg==", "dev": true }, "chokidar": { @@ -5091,9 +5079,9 @@ "dev": true }, "electron-to-chromium": { - "version": "1.5.286", - "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.286.tgz", - "integrity": "sha512-9tfDXhJ4RKFNerfjdCcZfufu49vg620741MNs26a9+bhLThdB+plgMeou98CAaHu/WATj2iHOOHTp1hWtABj2A==", + "version": "1.5.419", + "resolved": "https://registry.npmjs.org/electron-to-chromium/-/electron-to-chromium-1.5.419.tgz", + "integrity": "sha512-nHMPn8x4yCxCI0iSnL+LlHL5sUoUfjLXkcRIagZ4GBdrfFLFaiLNvzJWbJqZhFT9IAhw5tUSNlhggWN+otvp/A==", "dev": true }, "encodeurl": { @@ -5157,39 +5145,6 @@ "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", "dev": true }, - "eslint-scope": { - "version": "5.1.1", - "resolved": "https://registry.npmjs.org/eslint-scope/-/eslint-scope-5.1.1.tgz", - "integrity": "sha512-2NxwbF/hZ0KpepYN0cNbo+FN6XoK7GaHlQhgx/hIZl6Va0bF45RQOOwhLIy8lQDbuCiadSLCBnH2CFYquit5bw==", - "dev": true, - "requires": { - "esrecurse": "^4.3.0", - "estraverse": "^4.1.1" - } - }, - "esrecurse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/esrecurse/-/esrecurse-4.3.0.tgz", - "integrity": "sha512-KmfKL3b6G+RXvP8N1vr3Tq1kL/oCFgn2NYXEtqP8/L3pKapUA4G8cFVaoF3SU323CD4XypR/ffioHmkti6/Tag==", - "dev": true, - "requires": { - "estraverse": "^5.2.0" - }, - "dependencies": { - "estraverse": { - "version": "5.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-5.3.0.tgz", - "integrity": "sha512-MMdARuVEQziNTeJD8DgMqmhwR11BRQ/cBP+pLtYdSTnf3MIO8fFeiINEbX36ZdNlfU/7A9f3gUw49B3oQsvwBA==", - "dev": true - } - } - }, - "estraverse": { - "version": "4.3.0", - "resolved": "https://registry.npmjs.org/estraverse/-/estraverse-4.3.0.tgz", - "integrity": "sha512-39nnKffWz8xN1BU/2c79n9nB9HDzo0niYUqx6xyqUnyoAnQyyWpOTdZEeiCch8BBu515t4wp9ZmgVfVhn9EBpw==", - "dev": true - }, "etag": { "version": "1.8.1", "resolved": "https://registry.npmjs.org/etag/-/etag-1.8.1.tgz", @@ -5262,9 +5217,9 @@ "dev": true }, "fast-uri": { - "version": "3.1.5", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", - "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", + "version": "3.1.7", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.7.tgz", + "integrity": "sha512-dOvZVzjdZdz7phd9v6jCbwxrBW3fK6n8Rc0CtdmM4bumzMnxywBYhuph6J819RRw/ku+rLbelwfMunktuzVVHg==", "dev": true }, "fill-range": { @@ -5925,15 +5880,15 @@ } }, "minimizer-webpack-plugin": { - "version": "5.6.1", - "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.6.1.tgz", - "integrity": "sha512-DoeAZz8Q1C1znwsUzej1fdoi4jCf7/+Em27ouLqfK/+3m8G+D7yDhUwrc3CNhjSzGUN1kn7Iv4sWmjflQHenpw==", + "version": "5.8.0", + "resolved": "https://registry.npmjs.org/minimizer-webpack-plugin/-/minimizer-webpack-plugin-5.8.0.tgz", + "integrity": "sha512-2cT9+goJfBhtMz+gJqejSf09ClgmYhPhccRb/fb0ztVbixt0BkO8mRuI26FoPPuUNkRk6iEDryWAIouc5W8eJA==", "dev": true, "requires": { - "@jridgewell/trace-mapping": "^0.3.25", + "@jridgewell/trace-mapping": "^0.3.31", "jest-worker": "^27.4.5", - "schema-utils": "^4.3.0", - "terser": "^5.31.1" + "schema-utils": "^4.3.3", + "terser": "^5.51.0" } }, "ms": { @@ -5965,9 +5920,9 @@ "dev": true }, "node-releases": { - "version": "2.0.27", - "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.27.tgz", - "integrity": "sha512-nmh3lCkYZ3grZvqcCH+fjmQ7X+H0OeZgP40OierEaAptX4XofMh5kwNbWh7lBduUzCcV/8kZ+NDLCwm2iorIlA==", + "version": "2.0.54", + "resolved": "https://registry.npmjs.org/node-releases/-/node-releases-2.0.54.tgz", + "integrity": "sha512-YHs7BmmcsdAI5Ozuf8JZo6PT0mv2GIWC9vMfvUC3dp65M8hn7Ux8CPL+2oBI7juNuj9d0ndhTcznq2ODBps9cQ==", "dev": true }, "normalize-path": { @@ -6550,9 +6505,9 @@ "dev": true }, "terser": { - "version": "5.50.0", - "resolved": "https://registry.npmjs.org/terser/-/terser-5.50.0.tgz", - "integrity": "sha512-CN9BVxWhgS/hRxtUMjtC2uRWSTcSfQFHMDWma6sKKfIivCD91sM+FOPfvwoaRMqCSrUpe1nv3jDamd9eEQ4y+w==", + "version": "5.51.2", + "resolved": "https://registry.npmjs.org/terser/-/terser-5.51.2.tgz", + "integrity": "sha512-bWnjSNscmuI+GJze6ZupnHP8G/cTcsJF+bXCeQknk2SHQsgbNJnLrqiH9jZ2W4STPVXH2mDKKRX3iwPhc9Cn/Q==", "dev": true, "requires": { "@jridgewell/source-map": "^0.3.3", @@ -6663,9 +6618,9 @@ "dev": true }, "update-browserslist-db": { - "version": "1.2.3", - "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.2.3.tgz", - "integrity": "sha512-Js0m9cx+qOgDxo0eMiFGEueWztz+d4+M3rGlmKPT+T4IS/jP4ylw3Nwpu6cpTTP8R1MAC1kF4VbdLt3ARf209w==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/update-browserslist-db/-/update-browserslist-db-1.3.2.tgz", + "integrity": "sha512-UQ+MSxlhRm1bzjhU+DcuXfjFO1FzNtqhK5+9Yvlp90ItDLk5vT932A0rFu619nf7RVS+Y/VeaUW1jaRDqZ8VJw==", "dev": true, "requires": { "escalade": "^3.2.0", @@ -6688,9 +6643,9 @@ } }, "webpack": { - "version": "5.109.2", - "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.109.2.tgz", - "integrity": "sha512-U9/cvLzxObKNEZ9+TtdqrHM5/9z3lgl2c+c4BzbqGxFQvQvBAq87yql5A8pQ+rrMbS496MZJeF5enVBndIy2hw==", + "version": "5.110.1", + "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.110.1.tgz", + "integrity": "sha512-gInQB+jxXxgnZyvPwuzT5NGQmECDqeu85oxcrjinrYHqPoBex0hCAN2SFTJVyPVrK0Pq9E44VFP+e89fAc10/w==", "dev": true, "requires": { "@types/estree": "^1.0.8", @@ -6703,11 +6658,10 @@ "chrome-trace-event": "^1.0.2", "enhanced-resolve": "^5.24.4", "es-module-lexer": "^2.1.0", - "eslint-scope": "5.1.1", "events": "^3.2.0", "graceful-fs": "^4.2.11", "mime-db": "^1.54.0", - "minimizer-webpack-plugin": "^5.6.1", + "minimizer-webpack-plugin": "^5.7.0", "neo-async": "^2.6.2", "schema-utils": "^4.3.3", "tapable": "^2.3.0", @@ -6724,9 +6678,9 @@ } }, "webpack-cli": { - "version": "7.2.2", - "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.2.tgz", - "integrity": "sha512-lD0pALneslq8FfV+rwvm1BMW0AFAJrHHhNupAGN4asYjMvqrtRsenU4iKpiBo09gS4ntMxKGUxl9jhTEzVt0oA==", + "version": "7.2.3", + "resolved": "https://registry.npmjs.org/webpack-cli/-/webpack-cli-7.2.3.tgz", + "integrity": "sha512-vDFU7jrfCctnN7jJQWPl+V26B51GLp11prVZXg50oeonsgeBzSTJEWmXkjHsOjTgMjlPOQM8GWLh33X9RN/0ow==", "dev": true, "requires": { "@discoveryjs/json-ext": "^1.1.0", diff --git a/datafusion/wasmtest/datafusion-wasm-app/package.json b/datafusion/wasmtest/datafusion-wasm-app/package.json index c5723b079b6d0..ab03cbf969b81 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package.json @@ -27,8 +27,8 @@ "datafusion-wasmtest": "../pkg" }, "devDependencies": { - "webpack": "5.109.2", - "webpack-cli": "7.2.2", + "webpack": "5.110.1", + "webpack-cli": "7.2.3", "webpack-dev-server": "6.0.0", "copy-webpack-plugin": "14.0.0" } diff --git a/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md b/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md index d81d9eb9adab5..5366080ffe051 100644 --- a/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md +++ b/docs/source/contributor-guide/gsoc/gsoc_project_ideas_2025.md @@ -2,7 +2,7 @@ ## Introduction -Welcome to the Apache DataFusion Google Summer of Code (GSoC) 2025 project ideas list. Below you can find information about the projects. Please refer to [this page](https://datafusion.apache.org/contributor-guide/gsoc_application_guidelines.html) for application guidelines. +Welcome to the Apache DataFusion Google Summer of Code (GSoC) 2025 project ideas list. Below you can find information about the projects. Please refer to [this page](https://datafusion.apache.org/contributor-guide/gsoc/gsoc_application_guidelines_2025.html) for application guidelines. ## Projects diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 97ac11a0169a0..881aee0c6308f 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -97,4 +97,25 @@ The output type of the `floor` and `ceil` UDFs has been changed from the exact i Change the expected type or wrap the expression in `CAST`. It's recommended to avoid relying on decimal's exact precision and scale. -[#24703]: https://github.com/apache/datafusion/pull/24703 +### `DataFrame::from_columns` accepts `IntoIterator` + +`DataFrame::from_columns` now accepts any `IntoIterator` +instead of specifically accepting a `Vec<(&str, ArrayRef)>`. + +```rust,ignore +// Existing Vec usage continues to work +let df = DataFrame::from_columns(vec![ + ("id", id), + ("name", name), +])?; + +// Arrays can now be used directly +let df = DataFrame::from_columns([ + ("id", id), + ("name", name), +])?; +``` + +Most existing call sites using `Vec` require no changes. Code that relies on +the exact non-generic function signature of `DataFrame::from_columns` may need +to be updated to account for the new generic API. diff --git a/docs/source/user-guide/concepts-readings-events.md b/docs/source/user-guide/concepts-readings-events.md index 366defa35eb5f..9edd83b33a4b7 100644 --- a/docs/source/user-guide/concepts-readings-events.md +++ b/docs/source/user-guide/concepts-readings-events.md @@ -53,13 +53,13 @@ This is a list of DataFusion related blog posts, articles, and other resources. - **2026-03-31** [Blog: Writing Custom Table Providers in Apache DataFusion](https://datafusion.apache.org/blog/2026/03/31/writing-table-providers/) -- **2026-03-20** [Blog: Turning LIMIT into an I/O Optimization: Inside DataFusion’s Multi-Layer Pruning Stack](https://datafusion.apache.org/blog/2026/03/20/multi-layer-pruning/) +- **2026-03-20** [Blog: Turning LIMIT into an I/O Optimization: Inside DataFusion’s Multi-Layer Pruning Stack](https://datafusion.apache.org/blog/2026/03/20/limit-pruning/) - **2026-02-23** [Blog: Apache DataFusion: A Data Engineer's Guide to the Query Engine Reshaping How We Build Data Systems](https://andrewmadson.substack.com/p/apache-datafusion-a-data-engineers) - **2026-02-09** [Blog: Vector search using only Parquet and DataFusion](https://blog.xiangpeng.systems/posts/vector-search-with-parquet-datafusion/) -- **2026-02-02** [Blog: Optimizing SQL CASE Expression Evaluation](https://datafusion.apache.org/blog/2026/02/02/case-expression/) +- **2026-02-02** [Blog: Optimizing SQL CASE Expression Evaluation](https://datafusion.apache.org/blog/2026/02/02/datafusion_case/) - **2026-01-12** [Blog: Extending SQL in DataFusion: from ->> to TABLESAMPLE](https://datafusion.apache.org/blog/2026/01/12/extending-sql)