perf: backport bitmap pruning for runtime join filters
diff --git a/native/Cargo.lock b/native/Cargo.lock index c73688f..81dd9ac 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock
@@ -100,6 +100,15 @@ ] [[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] name = "allocator-api2" version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -243,6 +252,8 @@ "arrow-schema", "arrow-select", "arrow-string", + "half", + "rand 0.9.5", ] [[package]] @@ -1242,6 +1253,16 @@ ] [[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", + "serde_core", +] + +[[package]] name = "bumpalo" version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1544,6 +1565,17 @@ ] [[package]] +name = "console" +version = "0.16.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e96a4956774c13c126a8b5af4daa79384f4d826534c95a02d76afb39e2ab64e3" +dependencies = [ + "encode_unicode", + "libc", + "windows-sys 0.61.2", +] + +[[package]] name = "const-oid" version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1663,7 +1695,7 @@ "cast", "ciborium", "clap", - "criterion-plot", + "criterion-plot 0.6.0", "itertools 0.13.0", "num-traits", "oorandom", @@ -1678,6 +1710,32 @@ ] [[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot 0.8.2", + "futures", + "itertools 0.13.0", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] name = "criterion-plot" version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1688,6 +1746,16 @@ ] [[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools 0.13.0", +] + +[[package]] name = "crossbeam-channel" version = "0.5.17" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1997,7 +2065,7 @@ "comet-test-udfs", "comet-udf-sdk", "crc32fast", - "criterion", + "criterion 0.7.0", "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", @@ -2085,7 +2153,7 @@ version = "1.2.0" dependencies = [ "arrow", - "criterion", + "criterion 0.7.0", "datafusion", "datafusion-comet-common", "datafusion-comet-spark-expr", @@ -2114,7 +2182,7 @@ "clap", "crc32c", "crc32fast", - "criterion", + "criterion 0.7.0", "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", @@ -2139,7 +2207,7 @@ "arrow", "base64 0.23.1", "chrono", - "criterion", + "criterion 0.7.0", "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", @@ -2581,6 +2649,7 @@ "datafusion-expr-common", "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", + "datafusion-proto-models", "half", "hashbrown 0.17.1", "indexmap 2.14.2", @@ -2615,6 +2684,7 @@ "chrono", "datafusion-common", "datafusion-expr-common", + "datafusion-proto-models", "hashbrown 0.17.1", "indexmap 2.14.2", "itertools 0.15.0", @@ -2644,8 +2714,6 @@ [[package]] name = "datafusion-physical-plan" version = "55.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52" dependencies = [ "arrow", "arrow-data", @@ -2654,42 +2722,76 @@ "arrow-schema", "async-trait", "bytes", + "criterion 0.8.2", "datafusion-common", "datafusion-common-runtime", "datafusion-execution", "datafusion-expr", "datafusion-functions", + "datafusion-functions-aggregate", "datafusion-functions-aggregate-common", + "datafusion-functions-window", "datafusion-functions-window-common", "datafusion-physical-expr", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "futures", "half", "hashbrown 0.17.1", "indexmap 2.14.2", + "insta", "itertools 0.15.0", "log", "num-traits", "parking_lot", "pin-project-lite", + "rand 0.9.5", + "rstest", + "rstest_reuse", "serde_json", "tokio", ] [[package]] -name = "datafusion-pruning" +name = "datafusion-proto-common" version = "55.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492" +checksum = "57901d70c724dab7cb604fea579f5324332244e3c400ff2d7c9b5fc065051437" +dependencies = [ + "arrow", + "datafusion-common", + "prost", +] + +[[package]] +name = "datafusion-proto-models" +version = "55.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a574908455152a11abea49c7916cc8a504de946e5302c7e40bab632cbbcb9b2" +dependencies = [ + "datafusion-common", + "datafusion-proto-common", + "prost", +] + +[[package]] +name = "datafusion-pruning" +version = "55.1.0" dependencies = [ "arrow", "datafusion-common", "datafusion-datasource", + "datafusion-expr", "datafusion-expr-common", + "datafusion-functions-nested", "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "insta", + "itertools 0.15.0", "log", + "rstest", ] [[package]] @@ -2943,6 +3045,12 @@ checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" [[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] name = "equator" version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3196,6 +3304,12 @@ checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" [[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + +[[package]] name = "futures-util" version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3284,6 +3398,19 @@ checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" [[package]] +name = "globset" +version = "0.4.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07c34a9410465b45bd9787443bc7370f37735bad04b0f0cd57ff1a3186c98988" +dependencies = [ + "aho-corasick", + "bstr", + "log", + "regex-automata", + "regex-syntax", +] + +[[package]] name = "gloo-timers" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3323,6 +3450,8 @@ "cfg-if", "crunchy", "num-traits", + "rand 0.9.5", + "rand_distr", "zerocopy", ] @@ -3828,6 +3957,22 @@ ] [[package]] +name = "insta" +version = "1.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" +dependencies = [ + "console", + "globset", + "once_cell", + "regex", + "similar", + "strip-ansi-escapes", + "tempfile", + "walkdir", +] + +[[package]] name = "inventory" version = "0.3.24" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -4881,6 +5026,16 @@ checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" [[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + +[[package]] name = "parking" version = "2.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -5288,6 +5443,15 @@ ] [[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + +[[package]] name = "proc-macro2" version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -5490,6 +5654,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" dependencies = [ + "libc", "rand_chacha 0.3.1", "rand_core 0.6.4", ] @@ -5560,6 +5725,16 @@ checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.5", +] + +[[package]] name = "rand_pcg" version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -5653,6 +5828,12 @@ checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + +[[package]] name = "reqsign-aliyun-oss" version = "3.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -5916,6 +6097,46 @@ ] [[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.119", + "unicode-ident", +] + +[[package]] +name = "rstest_reuse" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3a8fb4672e840a587a66fc577a5491375df51ddb88f2a2c2a792598c326fe14" +dependencies = [ + "quote", + "rand 0.8.8", + "syn 2.0.119", +] + +[[package]] name = "rust-ini" version = "0.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -6385,6 +6606,12 @@ checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" [[package]] +name = "similar" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" + +[[package]] name = "siphasher" version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -6477,6 +6704,15 @@ checksum = "7f446288b699d66d0fd2e30d1cfe7869194312524b3b9252594868ed26ef056a" [[package]] +name = "strip-ansi-escapes" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a8f8038e7e7969abb3f1b7c2a811225e9296da208539e0f79c5251d6cac0025" +dependencies = [ + "vte", +] + +[[package]] name = "strsim" version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -6587,7 +6823,7 @@ checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.2", @@ -6757,6 +6993,7 @@ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "signal-hook-registry", "socket2", @@ -6812,6 +7049,36 @@ ] [[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.15+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1340ea94a5856333492c9064b02c778b191dd2c853778d9609debdcdfea3a614" +dependencies = [ + "indexmap 2.14.2", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + +[[package]] name = "tower" version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -7083,6 +7350,15 @@ checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" [[package]] +name = "vte" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231fdcd7ef3037e8330d8e17e61011a2c244126acc0a982f4040ac3f9f0bc077" +dependencies = [ + "memchr", +] + +[[package]] name = "walkdir" version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -7465,6 +7741,15 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + +[[package]] name = "wit-bindgen" version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/native/Cargo.toml b/native/Cargo.toml index dac17e5..7ec6487 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml
@@ -17,7 +17,7 @@ [workspace] default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs"] -members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs"] +members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs", "vendor/datafusion-physical-plan", "vendor/datafusion-pruning"] # Crates under ../contrib are intentionally NOT workspace members. Core pulls them in as # optional path dependencies when their corresponding contrib feature is enabled. exclude = ["../contrib"] @@ -86,6 +86,12 @@ iceberg-storage-opendal = { git = "https://github.com/apache/iceberg-rust", rev = "af1da4c5bc86178c38c1c6db0578bdd9fb7021e3", features = ["opendal-memory", "opendal-fs", "opendal-s3", "opendal-gcs", "opendal-oss", "opendal-azdls"] } reqsign-core = "3" +# Backport DataFusion #25602 without changing Comet's Arrow/Iceberg dependencies. +# Keep both crates together; provenance and removal instructions are in vendor/README.md. +[patch.crates-io] +datafusion-physical-plan = { path = "vendor/datafusion-physical-plan" } +datafusion-pruning = { path = "vendor/datafusion-pruning" } + [profile.release] debug = true overflow-checks = false
diff --git a/native/core/src/execution/operators/dynamic_filter/join/tests.rs b/native/core/src/execution/operators/dynamic_filter/join/tests.rs index 14242b4..6ae974e 100644 --- a/native/core/src/execution/operators/dynamic_filter/join/tests.rs +++ b/native/core/src/execution/operators/dynamic_filter/join/tests.rs
@@ -1146,6 +1146,89 @@ assert_eq!(unclustered_enabled_filter_rows, 400); } +#[tokio::test] +async fn bitmap_filter_prunes_parquet_row_groups_inside_build_bounds() { + for mode in [PartitionMode::CollectLeft, PartitionMode::Partitioned] { + let mut scan_bytes = Vec::new(); + for enabled in [false, true] { + let mut config = SessionConfig::new() + .with_target_partitions(1) + .with_parquet_page_index_pruning(false); + // Isolate metadata pruning from filtering decoded rows or pages. + config.options_mut().execution.parquet.pushdown_filters = false; + let session = Arc::new(SessionContext::new_with_config(config)); + let (_file, scan) = parquet_probe((0..400).collect(), &session, 100); + let filter = filtered_probe(&scan); + let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let build = memory_exec( + [50, 350] + .into_iter() + .map(|key| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![key]))], + ) + .unwrap() + }) + .collect(), + ); + let join = + single_key_join_plans(build, Arc::clone(&filter) as Arc<dyn ExecutionPlan>, mode); + let plan = PhysicalPlanner::apply_join_dynamic_filter( + Arc::new(join), + enabled, + session.copied_config().options(), + ) + .unwrap(); + let expected = RecordBatch::try_new( + plan.schema(), + vec![ + Arc::new(Int32Array::from(vec![50, 350])), + Arc::new(Int32Array::from(vec![50, 350])), + ], + ) + .unwrap(); + let output = collect(Arc::clone(&plan), session.task_ctx()) + .await + .unwrap(); + assert_eq!( + batches_to_sort_string(&output), + batches_to_sort_string(&[expected]), + "{mode:?}, enabled={enabled}" + ); + + // Every row group overlaps the build's [50, 350] bounds. Only the + // bitmap can exclude the interior groups [100, 199] and [200, 299]. + // Comet's runtime join forces map membership rather than IN lists. + assert_eq!( + pruning_metric(&scan, "row_groups_pruned_statistics"), + if enabled { 2 } else { 0 }, + "{mode:?}, enabled={enabled}" + ); + assert_eq!( + filter.metrics().unwrap().output_rows().unwrap(), + if enabled { 200 } else { 400 } + ); + if enabled { + assert_eq!(metric(&plan, "dynamic_filter_join_filters_attached"), 1); + } + scan_bytes.push( + scan.metrics() + .unwrap() + .sum_by_name("bytes_scanned") + .unwrap() + .as_usize(), + ); + } + assert!( + scan_bytes[1] < scan_bytes[0], + "{mode:?}: bitmap pruning should avoid data reads: enabled={}, disabled={}", + scan_bytes[1], + scan_bytes[0] + ); + } +} + fn limited_session(bytes: usize) -> (SessionContext, Arc<dyn MemoryPool>) { let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(bytes)); let runtime = RuntimeEnvBuilder::new()
diff --git a/native/vendor/README.md b/native/vendor/README.md new file mode 100644 index 0000000..d7d3dad --- /dev/null +++ b/native/vendor/README.md
@@ -0,0 +1,91 @@ +<!-- +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. +--> + +# DataFusion bitmap pruning backport + +These two crates backport [DataFusion #25602](https://github.com/apache/datafusion/pull/25602) +to DataFusion 55.1.0. The native workspace patches both crates together, so hash joins and +Parquet pruning use the same bitmap type. All other DataFusion dependencies remain on the +existing locked release, as do Arrow, Parquet, and Iceberg. + +The upstream implementation first landed at +`4bb9f8d3192f6c945f2bf933231380e60341af66`, whose dependency graph uses Arrow 60. +Comet's pinned Iceberg revision uses Arrow 59. This temporary backport makes the bitmap +available without also requiring an Iceberg migration or an unmerged Iceberg dependency. +It does not enable Comet's experimental join runtime filtering by default or add Delta scans. + +## Provenance + +The source archives are the published crates.io 55.1.0 packages, based on DataFusion release +commit `7d3835c71f30cbd3c3ae4041732267f1f453097a`. Their SHA-256 checksums match the +original entries in Comet's lockfile: + +| Package | Archive SHA-256 | +| -------------------------- | ------------------------------------------------------------------ | +| `datafusion-physical-plan` | `1265d58e5bce07d154e642a51ff43033b576a6ae40d50a29ac2b9f311004eb52` | +| `datafusion-pruning` | `74b06b333405c05015836ed36cea02cc29d1e2ee8fb414a77532d8346fb56492` | + +The crates retain their licenses, notices, normalized Cargo manifests, and original manifests. +Manifest license headers and the upstream Rust formatting configuration are included for +Comet's repository checks. Deprecation warnings are allowed in the physical-plan crate so +its pre-Rust-1.99 atomic API remains compatible with DataFusion 55.1.0's minimum Rust version. +Package-local lockfiles and Cargo cache markers are omitted; `native/Cargo.lock` controls the +build. [bitmap-backport.patch](bitmap-backport.patch) records every source and manifest change +relative to those archives. + +## Adaptations to 55.1.0 + +- Preserve the four-argument `HashTableLookupExpr::new` API. The new + `with_pruning_bitmap` builder attaches the optional summary without changing other callers. +- Use 55.1.0's existing hash-join build path. It has no prepared-build API; skip the summary + whenever dynamic filters are not computed. Failure to reserve bitmap memory only disables + this optional optimization. +- Retain the merged PR's conservative handling of `CASE` without `ELSE`, which must preserve + implicit NULLs during Parquet full-match inference. +- Keep existing IN-list pruning unchanged. Port only the bitmap statistics rewrite and CASE + support, adapting the recursive call to the 55.1.0 signature. Recognize lossless casts from + unsigned integers to larger signed integers locally, since 55.1.0's cast helper predates + that support. +- Include upstream bitmap, expression-sharing, and memory-limit tests, plus regressions for + missing statistics and implicit NULLs. Comet's join test verifies that both supported native + hash-join modes skip interior row groups that min/max filtering alone cannot exclude. + +## Validation + +From `native/`, run: + +```sh +cargo test -p datafusion-pruning --lib +cargo test -p datafusion-physical-plan --lib key_range_bitmap +cargo test -p datafusion-physical-plan --lib pruning_bitmap +cargo test -p datafusion-comet --lib dynamic_filter +``` + +The Comet regression is named `bitmap_filter_prunes_parquet_row_groups_inside_build_bounds`. +It checks exact join results, reader attachment, row groups pruned, decoded rows, and bytes +read with filtering enabled and disabled. Page and decode-time row filtering are disabled +in this test to isolate row-group pruning. + +## Removal + +When Comet can adopt a DataFusion release containing #25602 together with compatible Arrow +and Iceberg dependencies, remove both vendored crates, their workspace-member entries, and +the `[patch.crates-io]` entries in `native/Cargo.toml`, then update the lockfile. Retain the +Comet regression test. Do not update just one patched crate: the producer and pruning reader +must agree on the hash-lookup expression and bitmap implementation.
diff --git a/native/vendor/bitmap-backport.patch b/native/vendor/bitmap-backport.patch new file mode 100644 index 0000000..d55725c --- /dev/null +++ b/native/vendor/bitmap-backport.patch
@@ -0,0 +1,1175 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +--- a/datafusion-physical-plan/Cargo.toml ++++ b/datafusion-physical-plan/Cargo.toml +@@ -1,3 +1,20 @@ ++# Licensed to the Apache Software Foundation (ASF) under one ++# or more contributor license agreements. See the NOTICE file ++# distributed with this work for additional information ++# regarding copyright ownership. The ASF licenses this file ++# to you under the Apache License, Version 2.0 (the ++# "License"); you may not use this file except in compliance ++# with the License. You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, ++# software distributed under the License is distributed on an ++# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY ++# KIND, either express or implied. See the License for the ++# specific language governing permissions and limitations ++# under the License. ++ + # THIS FILE IS AUTOMATICALLY GENERATED BY CARGO + # + # When uploading crates to the registry Cargo will automatically +@@ -276,6 +293,9 @@ + used_underscore_binding = "warn" + + [lints.rust] ++# Rust 1.99 renamed AtomicUsize::fetch_update. Retain the upstream implementation ++# and its Rust 1.94 minimum instead of requiring the new try_update method. ++deprecated = "allow" + unused_qualifications = "deny" + + [lints.rust.unexpected_cfgs] +--- a/datafusion-physical-plan/src/joins/hash_join/exec.rs ++++ b/datafusion-physical-plan/src/joins/hash_join/exec.rs +@@ -40,6 +40,7 @@ + BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, + }; + use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64}; ++use crate::joins::key_range_bitmap::KeyRangeBitmap; + use crate::joins::utils::{ + OnceAsync, OnceFut, asymmetric_join_output_partitioning, reorder_output_after_swap, + swap_join_projection, update_hash, +@@ -2384,7 +2385,8 @@ + + let map = Arc::new(join_hash_map); + +- let membership = if num_rows == 0 { ++ // Nothing reads the strategy unless this join computes dynamic filters. ++ let membership = if num_rows == 0 || !should_compute_dynamic_filters { + PushdownStrategy::Empty + } else { + // If the build side is small enough we can use IN list pushdown. +@@ -2394,19 +2396,45 @@ + .iter() + .map(|arr| arr.get_array_memory_size()) + .sum::<usize>(); +- if left_values.is_empty() +- || left_values[0].is_empty() +- || estimated_size > config.optimizer.hash_join_inlist_pushdown_max_size +- || map.num_of_distinct_key() +- > config ++ ++ let pushdown_inlist = !left_values.is_empty() ++ && !left_values[0].is_empty() ++ && estimated_size <= config.optimizer.hash_join_inlist_pushdown_max_size ++ && map.num_of_distinct_key() ++ <= config + .optimizer +- .hash_join_inlist_pushdown_max_distinct_values ++ .hash_join_inlist_pushdown_max_distinct_values; ++ ++ if pushdown_inlist ++ && let Some(in_list_values) = build_struct_inlist_values(&left_values)? + { +- PushdownStrategy::Map(Arc::clone(&map)) +- } else if let Some(in_list_values) = build_struct_inlist_values(&left_values)? { + PushdownStrategy::InList(in_list_values) + } else { +- PushdownStrategy::Map(Arc::clone(&map)) ++ // Past the InList threshold use a bucket bitmap for container pruning. ++ let pruning_bitmap = match (left_values.as_slice(), bounds.as_ref()) { ++ ([keys], Some(bounds)) if !keys.is_empty() => bounds ++ .get_column_bounds(0) ++ .and_then(|b| { ++ KeyRangeBitmap::try_new( ++ keys, ++ &b.min, ++ &b.max, ++ map.num_of_distinct_key(), ++ ) ++ }) ++ .map(Arc::new), ++ _ => None, ++ }; ++ // Held for the join's lifetime, so charge it like the maps; it is ++ // optional, so skip it rather than fail when the pool is full. ++ let pruning_bitmap = pruning_bitmap.filter(|bitmap| { ++ let ok = reservation.try_grow(bitmap.size()).is_ok(); ++ if ok { ++ metrics.build_mem_used.add(bitmap.size()); ++ } ++ ok ++ }); ++ PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap) + } + }; + +@@ -2503,6 +2531,71 @@ + use insta::{allow_duplicates, assert_snapshot}; + use rstest::*; + use rstest_reuse::*; ++ ++ /// Runs a join whose 200 keys spread over 2M would size a pruning bitmap ++ /// at the 128 KiB cap, and reports the bytes the build side reserved. ++ /// ++ /// Forces the IN-list threshold off: under `force_hash_collisions` every ++ /// key hashes to the same bucket, so `num_of_distinct_key()` no longer ++ /// reflects the real 200 distinct keys and would otherwise take the ++ /// IN-list branch instead of the one under test. ++ async fn build_mem_used(limit: usize, dynamic_filters: bool) -> Result<usize> { ++ let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); ++ let batch = |keys: Vec<i64>| { ++ let column = Arc::new(Int64Array::from(keys)) as ArrayRef; ++ let batch = RecordBatch::try_new(Arc::clone(&schema), vec![column])?; ++ TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) ++ }; ++ let on = vec![( ++ Arc::new(Column::new_with_schema("k", &schema)?) as _, ++ Arc::new(Column::new_with_schema("k", &schema)?) as _, ++ )]; ++ let build = batch((0..200).map(|i| i * 10_000).collect())?; ++ let probe = batch(vec![0, 500_000, 1_500_000])?; ++ let join = if dynamic_filters { ++ hash_join_with_dynamic_filter(build, probe, on, JoinType::Inner)?.0 ++ } else { ++ join( ++ build, ++ probe, ++ on, ++ &JoinType::Inner, ++ NullEquality::NullEqualsNothing, ++ )? ++ }; ++ ++ let runtime = RuntimeEnvBuilder::new() ++ .with_memory_limit(limit, 1.0) ++ .build_arc()?; ++ let mut config = SessionConfig::new(); ++ config ++ .options_mut() ++ .optimizer ++ .hash_join_inlist_pushdown_max_size = 0; ++ let task_ctx = Arc::new( ++ TaskContext::default() ++ .with_runtime(runtime) ++ .with_session_config(config), ++ ); ++ let batches = common::collect(join.execute(0, task_ctx)?).await?; ++ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3); ++ Ok(join ++ .metrics() ++ .unwrap() ++ .sum_by_name("build_mem_used") ++ .unwrap() ++ .as_usize()) ++ } ++ ++ /// The bitmap is pruning-only: never built when no dynamic filter will read ++ /// it, and dropped rather than fatal when the pool cannot fit it. ++ #[tokio::test] ++ async fn pruning_bitmap_is_optional() -> Result<()> { ++ assert!(build_mem_used(1_000_000, false).await? < 100_000); ++ assert!(build_mem_used(1_000_000, true).await? > 100_000); ++ build_mem_used(100_000, true).await?; ++ Ok(()) ++ } + + #[derive(Debug)] + struct PartitionedTestExec { +--- a/datafusion-physical-plan/src/joins/hash_join/partitioned_hash_eval.rs ++++ b/datafusion-physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +@@ -35,6 +35,7 @@ + }; + + use crate::joins::Map; ++use crate::joins::key_range_bitmap::KeyRangeBitmap; + + /// RandomState wrapper that preserves the seed used to create it. + /// +@@ -265,6 +266,8 @@ + map: Arc<Map>, + /// Description for display + description: String, ++ /// Shared pruning-only view of the build side's key range ++ pruning_bitmap: Option<Arc<KeyRangeBitmap>>, + } + impl HashTableLookupExpr { + /// Create a new HashTableLookupExpr +@@ -288,7 +291,19 @@ + random_state, + map, + description, +- } ++ pruning_bitmap: None, ++ } ++ } ++ ++ /// Attach a key-range summary for container pruning. ++ pub fn with_pruning_bitmap(mut self, bitmap: Option<Arc<KeyRangeBitmap>>) -> Self { ++ self.pruning_bitmap = bitmap; ++ self ++ } ++ ++ /// Which parts of this build side's key range hold keys, for container pruning. ++ pub fn pruning_bitmap(&self) -> Option<&Arc<KeyRangeBitmap>> { ++ self.pruning_bitmap.as_ref() + } + } + impl std::fmt::Debug for HashTableLookupExpr { +@@ -353,12 +368,13 @@ + self: Arc<Self>, + children: Vec<Arc<dyn PhysicalExpr>>, + ) -> Result<Arc<dyn PhysicalExpr>> { +- Ok(Arc::new(HashTableLookupExpr::new( +- children, +- self.random_state.clone(), +- Arc::clone(&self.map), +- self.description.clone(), +- ))) ++ Ok(Arc::new(HashTableLookupExpr { ++ on_columns: children, ++ random_state: self.random_state.clone(), ++ map: Arc::clone(&self.map), ++ description: self.description.clone(), ++ pruning_bitmap: self.pruning_bitmap.clone(), ++ })) + } + + fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { +@@ -438,6 +454,7 @@ + mod tests { + use super::*; + use crate::joins::join_hash_map::JoinHashMapU32; ++ use datafusion_common::ScalarValue; + use datafusion_physical_expr::expressions::Column; + use std::collections::hash_map::DefaultHasher; + use std::hash::Hasher; +@@ -446,6 +463,38 @@ + let mut hasher = DefaultHasher::new(); + value.hash(&mut hasher); + hasher.finish() ++ } ++ ++ /// The pruning bitmap must survive `with_new_children`, which a scan calls ++ /// once per file. ++ #[test] ++ fn test_pruning_bitmap_shared_with_derived_children() { ++ let keys: ArrayRef = Arc::new(arrow::array::Int64Array::from(vec![0, 1_000_000])); ++ let bitmap = Arc::new( ++ KeyRangeBitmap::try_new( ++ &keys, ++ &ScalarValue::Int64(Some(0)), ++ &ScalarValue::Int64(Some(1_000_000)), ++ 2, ++ ) ++ .expect("two keys a million apart leave prunable gaps"), ++ ); ++ let expr = Arc::new( ++ HashTableLookupExpr::new( ++ vec![Arc::new(Column::new("a", 0))], ++ SeededRandomState::with_seed(1), ++ Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))), ++ "hash_lookup".to_string(), ++ ) ++ .with_pruning_bitmap(Some(Arc::clone(&bitmap))), ++ ); ++ ++ let derived = expr ++ .with_new_children(vec![Arc::new(Column::new("a", 7))]) ++ .unwrap(); ++ let derived = derived.downcast_ref::<HashTableLookupExpr>().unwrap(); ++ assert_eq!(derived.children()[0].to_string(), "a@7"); ++ assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap())); + } + + #[test] +--- a/datafusion-physical-plan/src/joins/hash_join/shared_bounds.rs ++++ b/datafusion-physical-plan/src/joins/hash_join/shared_bounds.rs +@@ -31,6 +31,7 @@ + use crate::joins::hash_join::partitioned_hash_eval::{ + HashExpr, HashTableLookupExpr, SeededRandomState, + }; ++use crate::joins::key_range_bitmap::KeyRangeBitmap; + use crate::repartition::RangeExpr; + use arrow::array::ArrayRef; + use arrow::datatypes::{DataType, Field, Schema}; +@@ -136,12 +137,16 @@ + )?))) + } + // Use hash table lookup for large build sides +- PushdownStrategy::Map(hash_map) => Ok(Some(Arc::new(HashTableLookupExpr::new( +- on_right.to_vec(), +- random_state.clone(), +- hash_map, +- "hash_lookup".to_string(), +- )) as Arc<dyn PhysicalExpr>)), ++ PushdownStrategy::Map(hash_map, pruning_bitmap) => Ok(Some(Arc::new( ++ HashTableLookupExpr::new( ++ on_right.to_vec(), ++ random_state.clone(), ++ hash_map, ++ "hash_lookup".to_string(), ++ ) ++ .with_pruning_bitmap(pruning_bitmap), ++ ) ++ as Arc<dyn PhysicalExpr>)), + // Empty partition - should not create a filter for this + PushdownStrategy::Empty => Ok(None), + } +@@ -277,8 +282,9 @@ + pub(crate) enum PushdownStrategy { + /// Use InList for small build sides (< 128MB) + InList(ArrayRef), +- /// Use map lookup for large build sides +- Map(Arc<Map>), ++ /// Use map lookup for large build sides. If `Some` the second field ++ /// represents the build side's keys using a bucket bitmap. ++ Map(Arc<Map>, Option<Arc<KeyRangeBitmap>>), + /// There was no data in this partition, do not build a dynamic filter for it + Empty, + } +--- /dev/null ++++ b/datafusion-physical-plan/src/joins/key_range_bitmap.rs +@@ -0,0 +1,238 @@ ++// 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. ++ ++//! Bitmap over a hash join build side's key range, for container pruning. ++ ++use arrow::array::{Array, ArrayRef, BooleanArray, downcast_integer_array}; ++use arrow::buffer::{BooleanBuffer, MutableBuffer}; ++use arrow::util::bit_util; ++use datafusion_common::ScalarValue; ++use num_traits::AsPrimitive; ++ ++use crate::joins::array_map::ArrayMap; ++ ++/// Largest bitmap, in bits (128 KiB), which bounds what one build side can hold. ++const MAX_BUCKETS: u64 = 1 << 20; ++ ++/// Which parts of a build side's key range hold at least one key. ++/// ++/// Keys map order-preservingly onto buckets spanning `[min, max]`, one bit each, ++/// so "can a container holding `[lo, hi]` match?" is a scan of the bits spanning ++/// `[lo, hi]`: all clear proves no build key lies in it. ++/// ++/// Below, keys 3 to 97 over eight buckets: `offset` is 3, `span` is 97 - 3, and ++/// eight buckets need `shift` = 4, so each covers `1 << 4` keys starting at ++/// `offset`. Only buckets 0 to `span >> shift` can hold a key. Real bitmaps run ++/// to `MAX_BUCKETS`; eight is what fits on a line. ++/// ++/// ```text ++/// keys 3 12 42 97 ++/// v v v v ++/// bucket | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | ++/// covers 3-18 19-34 35-50 51-66 67-82 83-98 unused unused ++/// bit 1 0 1 0 0 1 0 0 ++/// ``` ++/// ++/// A container is kept when any bucket its `[min, max]` touches has a bit set. ++#[derive(Debug)] ++pub struct KeyRangeBitmap { ++ /// Smallest key, in the `u64` ordering the join's array map uses. ++ offset: u64, ++ /// `max - min`: a key is in range iff `key.wrapping_sub(offset) <= span`. ++ span: u64, ++ /// Each bucket covers `1 << shift` keys. ++ shift: u32, ++ /// One bit per bucket. ++ bits: BooleanBuffer, ++} ++ ++impl KeyRangeBitmap { ++ /// Maps `array`'s non-null values, which must lie within `[min, max]`. ++ /// ++ /// `None` when the key type has no `u64` ordering, or when the bitmap could ++ /// not exclude anything the `[min, max]` bounds predicate does not already. ++ pub fn try_new( ++ array: &ArrayRef, ++ min: &ScalarValue, ++ max: &ScalarValue, ++ distinct_keys: usize, ++ ) -> Option<Self> { ++ if !ArrayMap::is_supported_type(array.data_type()) { ++ return None; ++ } ++ let offset = ArrayMap::key_to_u64(min)?; ++ let span = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); ++ ++ // One bucket per key value where the range allows it, so the bitmap is ++ // exact; otherwise as many as the cap permits. ++ let buckets = span.saturating_add(1).min(MAX_BUCKETS).next_power_of_two(); ++ ++ // Use a heuristic to bail-out for the common shape which renders each bucket ++ // set (contiguous surrogate keys, a date range) before actually computing it. ++ if (distinct_keys as u64).saturating_mul(2) > span.min(buckets) { ++ return None; ++ } ++ ++ let shift = (0u32..63).find(|s| (span >> s) < buckets).unwrap_or(63); ++ ++ let bits = bucket_bits(array, offset, shift, buckets); ++ ++ // Every bucket set means nothing can be excluded. ++ let reachable = ((span >> shift) + 1) as usize; ++ bits.slice(0, reachable).has_false().then_some(Self { ++ offset, ++ span, ++ shift, ++ bits, ++ }) ++ } ++ ++ /// One Boolean per `[min[i], max[i]]` container interval: `false` proves no ++ /// build key lies in that container, `None` where a bound is absent or the ++ /// statistics are not integer keys. ++ pub fn may_contain_ranges(&self, min: &dyn Array, max: &dyn Array) -> BooleanArray { ++ // Both bounds must be the same integer type for the tuple pattern to match. ++ downcast_integer_array!( ++ (min, max) => { ++ min.iter() ++ .zip(max.iter()) ++ .map(|bounds| match bounds { ++ (Some(lo), Some(hi)) => { ++ Some(self.may_contain_range(lo.as_(), hi.as_())) ++ } ++ _ => None, ++ }) ++ .collect() ++ } ++ _ => BooleanArray::new_null(min.len()), ++ ) ++ } ++ ++ /// Might a container spanning `[lo, hi]` hold a build key? `false` proves it ++ /// cannot. Bounds use the same `u64` key ordering as the join's array map. ++ pub fn may_contain_range(&self, lo: u64, hi: u64) -> bool { ++ // Offsets from the minimum; anything outside `[min, max]` wraps to ++ // something huge, which the clamping below turns into "past the end". ++ let lo_off = lo.wrapping_sub(self.offset); ++ let hi_off = hi.wrapping_sub(self.offset); ++ if lo_off > self.span && hi_off > self.span && lo_off <= hi_off { ++ return false; ++ } ++ let lo_b = if lo_off > self.span { ++ 0 ++ } else { ++ lo_off >> self.shift ++ }; ++ let hi_b = (hi_off.min(self.span) >> self.shift).max(lo_b); ++ self.bits ++ .slice(lo_b as usize, (hi_b - lo_b + 1) as usize) ++ .has_true() ++ } ++ ++ /// Bytes held, for the build-side memory reservation. ++ pub fn size(&self) -> usize { ++ self.bits.inner().capacity() ++ } ++ ++ /// Buckets set and buckets total, for `EXPLAIN`. ++ pub fn fill_stats(&self) -> (usize, usize) { ++ (self.bits.count_set_bits(), self.bits.len()) ++ } ++} ++ ++/// One bit per bucket, set for each bucket holding a non-null key of `keys`. ++fn bucket_bits(keys: &dyn Array, offset: u64, shift: u32, buckets: u64) -> BooleanBuffer { ++ let mut words = MutableBuffer::new_null(buckets as usize); ++ downcast_integer_array!( ++ keys => { ++ let bytes = words.as_slice_mut(); ++ let mut set = |key: u64| { ++ let bucket = (key.wrapping_sub(offset) >> shift).min(buckets - 1); ++ bit_util::set_bit(bytes, bucket as usize); ++ }; ++ if keys.null_count() == 0 { ++ keys.values().iter().for_each(|key| set((*key).as_())); ++ } else { ++ keys.iter().flatten().for_each(|key| set(key.as_())); ++ } ++ } ++ _ => unreachable!("guarded by ArrayMap::is_supported_type"), ++ ); ++ BooleanBuffer::new(words.into(), 0, buckets as usize) ++} ++ ++#[cfg(test)] ++mod tests { ++ use super::*; ++ use arrow::array::Int64Array; ++ use std::sync::Arc; ++ ++ fn bitmap(keys: &[i64]) -> Option<KeyRangeBitmap> { ++ let a: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); ++ KeyRangeBitmap::try_new( ++ &a, ++ &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), ++ &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), ++ keys.iter().collect::<std::collections::HashSet<_>>().len(), ++ ) ++ } ++ fn may(b: &KeyRangeBitmap, lo: i64, hi: i64) -> bool { ++ b.may_contain_range(lo as u64, hi as u64) ++ } ++ ++ #[test] ++ fn prunes_gaps_and_never_drops_a_match() { ++ let keys: Vec<i64> = (0..2_000_000).filter(|k| k % 10000 < 200).collect(); ++ let s = bitmap(&keys).expect("not saturated"); ++ let mut kept = 0; ++ for i in 0..2000i64 { ++ let (lo, hi) = (1000 * i, 1000 * i + 999); ++ let truth = keys ++ .binary_search(&lo) ++ .map_or_else(|p| keys.get(p).is_some_and(|k| *k <= hi), |_| true); ++ let got = may(&s, lo, hi); ++ assert!( ++ !truth || got, ++ "dropped a container holding a key: {lo}..{hi}" ++ ); ++ kept += got as usize; ++ } ++ assert!(kept <= 220, "expected ~200 kept, got {kept}"); ++ } ++ ++ #[test] ++ fn contiguous_keys_no_bitmap() { ++ assert!(bitmap(&(0..5000i64).collect::<Vec<_>>()).is_none()); ++ } ++ ++ #[test] ++ fn test_bitmap_range_probing() { ++ let s = bitmap(&[0, 1_000_000]).expect("not saturated"); ++ assert!(!may(&s, -500, -10)); ++ assert!(!may(&s, 2_000_000, 3_000_000)); ++ assert!(may(&s, -500, 10)); ++ ++ let s = bitmap(&[-1_000_000, -5, 0, 5, 1_000_000]).expect("not saturated"); ++ assert!(may(&s, -10, 10)); ++ assert!(!may(&s, -900_000, -800_000)); ++ ++ // The widest possible span must not overflow the bucket sizing. ++ let s = bitmap(&[i64::MIN, 0, i64::MAX]).expect("not saturated"); ++ assert!(may(&s, -1, 1)); ++ assert!(!may(&s, 1 << 50, 1 << 51)); ++ } ++} +--- a/datafusion-physical-plan/src/joins/mod.rs ++++ b/datafusion-physical-plan/src/joins/mod.rs +@@ -48,6 +48,7 @@ + /// Note: This module is public for internal testing purposes only + /// and is not guaranteed to be stable across versions. + pub mod join_hash_map; ++pub mod key_range_bitmap; + + use array_map::ArrayMap; + use utils::JoinHashMapType; +--- a/datafusion-pruning/Cargo.toml ++++ b/datafusion-pruning/Cargo.toml +@@ -1,3 +1,20 @@ ++# Licensed to the Apache Software Foundation (ASF) under one ++# or more contributor license agreements. See the NOTICE file ++# distributed with this work for additional information ++# regarding copyright ownership. The ASF licenses this file ++# to you under the Apache License, Version 2.0 (the ++# "License"); you may not use this file except in compliance ++# with the License. You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, ++# software distributed under the License is distributed on an ++# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY ++# KIND, either express or implied. See the License for the ++# specific language governing permissions and limitations ++# under the License. ++ + # THIS FILE IS AUTOMATICALLY GENERATED BY CARGO + # + # When uploading crates to the registry Cargo will automatically +@@ -81,6 +98,9 @@ + [dev-dependencies.itertools] + version = "0.15" + ++[dev-dependencies.rstest] ++version = "0.26.1" ++ + [lints.clippy] + allow_attributes = "warn" + assigning_clones = "warn" +--- a/datafusion-pruning/Cargo.toml.orig ++++ b/datafusion-pruning/Cargo.toml.orig +@@ -1,3 +1,20 @@ ++# Licensed to the Apache Software Foundation (ASF) under one ++# or more contributor license agreements. See the NOTICE file ++# distributed with this work for additional information ++# regarding copyright ownership. The ASF licenses this file ++# to you under the Apache License, Version 2.0 (the ++# "License"); you may not use this file except in compliance ++# with the License. You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, ++# software distributed under the License is distributed on an ++# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY ++# KIND, either express or implied. See the License for the ++# specific language governing permissions and limitations ++# under the License. ++ + [package] + name = "datafusion-pruning" + description = "DataFusion Pruning Logic" +--- /dev/null ++++ b/datafusion-pruning/src/key_range_bitmap_expr.rs +@@ -0,0 +1,103 @@ ++// Licensed to the Apache Software Foundation (ASF) under one ++// or more contributor license agreements. See the NOTICE file ++// distributed with this work for additional information ++// regarding copyright ownership. The ASF licenses this file ++// to you under the Apache License, Version 2.0 (the ++// "License"); you may not use this file except in compliance ++// with the License. You may obtain a copy of the License at ++// ++// http://www.apache.org/licenses/LICENSE-2.0 ++// ++// Unless required by applicable law or agreed to in writing, ++// software distributed under the License is distributed on an ++// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY ++// KIND, either express or implied. See the License for the ++// specific language governing permissions and limitations ++// under the License. ++ ++use std::fmt::{self, Display, Formatter}; ++use std::hash::{Hash, Hasher}; ++use std::sync::Arc; ++ ++use arrow::datatypes::{DataType, Schema}; ++use arrow::record_batch::RecordBatch; ++use datafusion_common::{Result, assert_eq_or_internal_err}; ++use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; ++use datafusion_physical_plan::ColumnarValue; ++use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; ++ ++/// Tests containers' `[min, max]` statistics against a build side's ++/// [`KeyRangeBitmap`]: one nullable Boolean per container, where `false` proves ++/// no build key lies in it and `NULL` means the bounds could not decide. ++#[derive(Debug)] ++pub(crate) struct KeyRangeBitmapPruningExpr { ++ pub(crate) min: PhysicalExprRef, ++ pub(crate) max: PhysicalExprRef, ++ pub(crate) bitmap: Arc<KeyRangeBitmap>, ++} ++ ++impl PartialEq for KeyRangeBitmapPruningExpr { ++ fn eq(&self, other: &Self) -> bool { ++ self.min.eq(&other.min) ++ && self.max.eq(&other.max) ++ && Arc::ptr_eq(&self.bitmap, &other.bitmap) ++ } ++} ++impl Eq for KeyRangeBitmapPruningExpr {} ++ ++impl Hash for KeyRangeBitmapPruningExpr { ++ fn hash<H: Hasher>(&self, state: &mut H) { ++ self.min.hash(state); ++ self.max.hash(state); ++ Arc::as_ptr(&self.bitmap).hash(state); ++ } ++} ++ ++impl Display for KeyRangeBitmapPruningExpr { ++ fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { ++ let (set, total) = self.bitmap.fill_stats(); ++ write!( ++ f, ++ "KEY_RANGE_BITMAP({}, {}, {set}/{total} buckets)", ++ self.min, self.max ++ ) ++ } ++} ++ ++impl PhysicalExpr for KeyRangeBitmapPruningExpr { ++ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { ++ Ok(DataType::Boolean) ++ } ++ ++ fn nullable(&self, _input_schema: &Schema) -> Result<bool> { ++ Ok(true) ++ } ++ ++ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { ++ let rows = batch.num_rows(); ++ let min = self.min.evaluate(batch)?.into_array(rows)?; ++ let max = self.max.evaluate(batch)?.into_array(rows)?; ++ let matches = self.bitmap.may_contain_ranges(&min, &max); ++ Ok(ColumnarValue::Array(Arc::new(matches))) ++ } ++ ++ fn children(&self) -> Vec<&PhysicalExprRef> { ++ vec![&self.min, &self.max] ++ } ++ ++ fn with_new_children( ++ self: Arc<Self>, ++ children: Vec<PhysicalExprRef>, ++ ) -> Result<PhysicalExprRef> { ++ assert_eq_or_internal_err!(children.len(), 2); ++ Ok(Arc::new(Self { ++ min: Arc::clone(&children[0]), ++ max: Arc::clone(&children[1]), ++ bitmap: Arc::clone(&self.bitmap), ++ })) ++ } ++ ++ fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { ++ write!(f, "{self}") ++ } ++} +--- a/datafusion-pruning/src/lib.rs ++++ b/datafusion-pruning/src/lib.rs +@@ -18,6 +18,7 @@ + #![cfg_attr(test, allow(clippy::needless_pass_by_value))] + + mod file_pruner; ++mod key_range_bitmap_expr; + mod pruning_predicate; + + pub use file_pruner::FilePruner; +--- a/datafusion-pruning/src/pruning_predicate.rs ++++ b/datafusion-pruning/src/pruning_predicate.rs +@@ -22,6 +22,8 @@ + use std::collections::HashSet; + use std::sync::Arc; + ++use crate::key_range_bitmap_expr::KeyRangeBitmapPruningExpr; ++ + use arrow::array::AsArray; + use arrow::{ + array::{ArrayRef, BooleanArray, new_null_array}, +@@ -48,6 +50,7 @@ + use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee}; + use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr}; + use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt; ++use datafusion_physical_plan::joins::HashTableLookupExpr; + use datafusion_physical_plan::{ColumnarValue, PhysicalExpr}; + + /// Used to prove that arbitrary predicates (boolean expression) can not +@@ -1466,6 +1469,89 @@ + /// `datafusion.execution.parquet.max_in_list_size` config option. + pub const MAX_IN_LIST_SIZE: usize = 20; + ++/// Statistics expressions for `column`: its min, its max, and "not every row is ++/// NULL". Appends to `required_columns`, rolling the appends back if any part ++/// cannot be built, since [`RequiredColumns::stat_column_expr`] only appends. ++fn min_max_non_null_exprs( ++ column: &phys_expr::Column, ++ column_expr: &Arc<dyn PhysicalExpr>, ++ field: &Field, ++ schema: &Schema, ++ required_columns: &mut RequiredColumns, ++) -> Option<(PhysicalExprRef, PhysicalExprRef, PhysicalExprRef)> { ++ let appended = required_columns.columns.len(); ++ let statistics = (|| { ++ let min = required_columns ++ .min_column_expr(column, column_expr, field) ++ .ok()?; ++ let max = required_columns ++ .max_column_expr(column, column_expr, field) ++ .ok()?; ++ // A cast on `column_expr` (widening it towards min/max's target type) ++ // does not affect nullability, so check it on the bare column beneath. ++ let bare_expr = match column_expr.downcast_ref::<phys_expr::CastExpr>() { ++ Some(cast) => cast.expr(), ++ None => column_expr, ++ }; ++ let non_null = ++ build_is_null_column_expr(bare_expr, schema, required_columns, true)?; ++ Some((min, max, non_null)) ++ })(); ++ if statistics.is_none() { ++ required_columns.columns.truncate(appended); ++ } ++ statistics ++} ++ ++/// Builds a "may match" expression for a [`HashTableLookupExpr`] (a large join ++/// build side pushed down as an opaque hash-table lookup): tests container ++/// min/max stats against the build side's key-range bitmap, which is fixed size ++/// and built once with the join's hash table. ++fn build_hash_lookup_pruning_expr( ++ lookup: &HashTableLookupExpr, ++ schema: &Schema, ++ required_columns: &mut RequiredColumns, ++) -> Option<Arc<dyn PhysicalExpr>> { ++ let bitmap = Arc::clone(lookup.pruning_bitmap()?); ++ let on_columns = lookup.children(); ++ let [column_expr] = on_columns[..] else { ++ return None; ++ }; ++ // The schema adapter often wraps the column in a same-type or widening cast ++ // (a nullable table column over a REQUIRED file column, or a narrower file ++ // type). `stat_column_expr` rewrites only the inner `Column`, keeping it. ++ let column = match column_expr.downcast_ref::<phys_expr::CastExpr>() { ++ Some(cast) => { ++ let column = cast.expr().downcast_ref::<phys_expr::Column>()?; ++ let from = schema.fields().get(column.index())?.data_type(); ++ // 55.1.0's widening-cast helper predates unsigned-to-larger-signed ++ // support. These additional integer casts preserve every value, ++ // ordering, and null; keep other unrecognized casts conservative. ++ let widens_unsigned = matches!( ++ (from, cast.cast_type()), ++ ( ++ DataType::UInt8, ++ DataType::Int16 | DataType::Int32 | DataType::Int64 ++ ) | (DataType::UInt16, DataType::Int32 | DataType::Int64) ++ | (DataType::UInt32, DataType::Int64) ++ ); ++ (cast.is_bigger_cast(from) || widens_unsigned).then_some(column)? ++ } ++ None => column_expr.downcast_ref::<phys_expr::Column>()?, ++ }; ++ let field = schema.fields().get(column.index())?; ++ if field.name() != column.name() { ++ return None; ++ } ++ let (min, max, non_null) = ++ min_max_non_null_exprs(column, column_expr, field, schema, required_columns)?; ++ Some(Arc::new(phys_expr::BinaryExpr::new( ++ non_null, ++ Operator::And, ++ Arc::new(KeyRangeBitmapPruningExpr { min, max, bitmap }), ++ ))) ++} ++ + /// Rewrite a predicate expression in terms of statistics (min/max/null_counts) + /// for use as a [`PruningPredicate`]. + pub struct PredicateRewriter { +@@ -1621,6 +1707,44 @@ + return unhandled_hook.handle(expr); + } + } ++ if let Some(lookup) = expr.downcast_ref::<HashTableLookupExpr>() { ++ return build_hash_lookup_pruning_expr(lookup, schema, required_columns) ++ .unwrap_or_else(|| unhandled_hook.handle(expr)); ++ } ++ // A partitioned hash join hides its per-partition filters under a `CASE` on the ++ // repartition hash. A row takes exactly one branch, so a container may match ++ // only if some branch may: the branches' disjunction is a sound relaxation, and ++ // the `WHEN`s (a hash, which no statistics describe) can be dropped. ++ if let Some(case) = expr.downcast_ref::<phys_expr::CaseExpr>() { ++ // Only a Boolean `CASE` is a predicate; anything else is a value for ++ // whatever compares it to handle. ++ if !matches!(case.data_type(schema), Ok(DataType::Boolean)) { ++ return unhandled_hook.handle(expr); ++ } ++ // Without an `ELSE`, unmatched rows are UNKNOWN: not a match, but not ++ // FALSE either, which full-match inference must be able to tell apart. ++ if case.else_expr().is_none() { ++ return unhandled_hook.handle(expr); ++ } ++ return case ++ .when_then_expr() ++ .iter() ++ .map(|(_, then)| then) ++ .chain(case.else_expr()) ++ .map(|branch| { ++ build_predicate_expression( ++ branch, ++ schema, ++ required_columns, ++ unhandled_hook, ++ max_in_list_size, ++ ) ++ }) ++ .reduce(|acc, branch| { ++ Arc::new(phys_expr::BinaryExpr::new(acc, Operator::Or, branch)) as _ ++ }) ++ .unwrap_or_else(|| unhandled_hook.handle(expr)); ++ } + + let (left, op, right) = { + if let Some(bin_expr) = expr.downcast_ref::<phys_expr::BinaryExpr>() { +@@ -2197,7 +2321,11 @@ + self as phys_expr, DynamicFilterPhysicalExpr, + }; + use datafusion_physical_expr::planner::logical2physical; ++ use datafusion_physical_plan::joins::join_hash_map::JoinHashMapU32; ++ use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; ++ use datafusion_physical_plan::joins::{Map, SeededRandomState}; + use itertools::Itertools; ++ use rstest::rstest; + + #[derive(Debug, Default)] + /// Mock statistic provider for tests +@@ -6061,4 +6189,236 @@ + "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1"; + assert_eq!(res.to_string(), expected); + } ++ ++ /// A hash-table lookup over `column` carrying a bitmap over `keys`. ++ fn hash_lookup( ++ column: &Arc<dyn PhysicalExpr>, ++ keys: &[i64], ++ ) -> Arc<dyn PhysicalExpr> { ++ let array: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); ++ let bitmap = KeyRangeBitmap::try_new( ++ &array, ++ &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), ++ &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), ++ keys.len(), ++ ); ++ Arc::new( ++ HashTableLookupExpr::new( ++ vec![Arc::clone(column)], ++ SeededRandomState::with_seed(1), ++ Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(1)))), ++ "hash_lookup".to_string(), ++ ) ++ .with_pruning_bitmap(bitmap.map(Arc::new)), ++ ) ++ } ++ ++ #[test] ++ fn test_hash_lookup_pruning_via_min_max() { ++ let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); ++ let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); ++ // Gapped keys: 10,20,30 then a far-away 100000, so the bitmap is sparse. ++ let lookup = hash_lookup(&column, &[10, 20, 30, 100_000]); ++ ++ let predicate = PruningPredicateBuilder::new() ++ .with_file_schema(Arc::clone(&schema)) ++ .try_build(lookup) ++ .unwrap(); ++ ++ // An opaque lookup yields no literal guarantees, so `contained()` is never ++ // consulted and every exclusion below comes from the min/max rewrite. ++ assert!(predicate.literal_guarantees().is_empty()); ++ ++ let statistics = TestStatistics::new().with( ++ "b", ++ ContainerStats::new_i64( ++ vec![Some(5000), Some(15), Some(50000)], ++ vec![Some(8000), Some(25), Some(60000)], ++ ), ++ ); ++ ++ let result = predicate.prune(&statistics).unwrap(); ++ // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 ++ // holds 20. The bitmap excludes the first and third. ++ assert_eq!(result, vec![false, true, false]); ++ } ++ ++ #[test] ++ fn test_hash_lookup_pruning_with_nulls_and_missing_statistics() { ++ let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); ++ let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); ++ let predicate = PruningPredicateBuilder::new() ++ .with_file_schema(schema) ++ .try_build(hash_lookup(&column, &[10, 20, 30, 100_000])) ++ .unwrap(); ++ let statistics = TestStatistics::new().with( ++ "b", ++ ContainerStats::new_i64( ++ vec![Some(15), None, None, None, Some(5000), Some(5000)], ++ vec![Some(25), None, None, Some(25), None, Some(8000)], ++ ) ++ .with_row_counts(vec![Some(2); 6]) ++ .with_null_counts(vec![ ++ Some(0), ++ Some(2), ++ Some(0), ++ Some(0), ++ Some(0), ++ None, ++ ]), ++ ); ++ ++ // A matching range is kept and an all-null group is rejected. Missing ++ // either bound cannot prove absence. A known gap still excludes a group ++ // when the null count is unknown: neither nulls nor non-nulls can match. ++ assert_eq!( ++ predicate.prune(&statistics).unwrap(), ++ vec![true, false, true, true, true, false] ++ ); ++ } ++ ++ #[test] ++ fn test_case_without_else_cannot_prove_full_match() { ++ let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); ++ let condition = logical2physical(&col("b").eq(lit(1)), &schema); ++ let case: Arc<dyn PhysicalExpr> = Arc::new( ++ phys_expr::CaseExpr::try_new( ++ None, ++ vec![(condition, phys_expr::lit(false))], ++ None, ++ ) ++ .unwrap(), ++ ); ++ // NOT(CASE WHEN b = 1 THEN false END) yields NULL for b=0, even ++ // though the referenced column is non-nullable. Such a row must not be ++ // counted as a match when a Parquet scan tries to satisfy a LIMIT. ++ let predicate: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::NotExpr::new(case)); ++ let batch = RecordBatch::try_new( ++ Arc::clone(&schema), ++ vec![Arc::new(Int32Array::from(vec![0, 1]))], ++ ) ++ .unwrap(); ++ assert_eq!( ++ predicate ++ .evaluate(&batch) ++ .unwrap() ++ .into_array(2) ++ .unwrap() ++ .as_boolean(), ++ &BooleanArray::from(vec![None, Some(true)]) ++ ); ++ ++ // Full-match inference negates the predicate, simplifies double NOT, ++ // and asks if the result prunes the group. Dropping the implicit NULL ++ // branch would produce false and incorrectly mark every row matched. ++ let inverse = PhysicalExprSimplifier::new(&schema) ++ .simplify(Arc::new(phys_expr::NotExpr::new(predicate))) ++ .unwrap(); ++ let inverse = PruningPredicateBuilder::new() ++ .with_file_schema(schema) ++ .try_build(inverse) ++ .unwrap(); ++ let statistics = TestStatistics::new().with( ++ "b", ++ ContainerStats::new_i32(vec![Some(0)], vec![Some(1)]) ++ .with_row_counts(vec![Some(2)]) ++ .with_null_counts(vec![Some(0)]), ++ ); ++ assert_eq!(inverse.prune(&statistics).unwrap(), vec![true]); ++ } ++ ++ #[rstest] ++ #[case::widening_i32_to_i64(DataType::Int32, DataType::Int64, true)] ++ #[case::widening_u32_to_i64(DataType::UInt32, DataType::Int64, true)] ++ #[case::widening_u16_to_i64(DataType::UInt16, DataType::Int64, true)] ++ #[case::same_type_i64_to_i64(DataType::Int64, DataType::Int64, true)] ++ #[case::narrowing_i64_to_i32(DataType::Int64, DataType::Int32, false)] ++ #[case::narrowing_u32_to_i16(DataType::UInt32, DataType::Int16, false)] ++ fn test_hash_lookup_pruning_through_cast( ++ #[case] file_type: DataType, ++ #[case] cast_type: DataType, ++ #[case] pruned: bool, ++ ) { ++ let schema = ++ Arc::new(Schema::new(vec![Field::new("b", file_type.clone(), true)])); ++ let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); ++ let cast_column: Arc<dyn PhysicalExpr> = ++ Arc::new(phys_expr::CastExpr::new(column, cast_type, None)); ++ let lookup = hash_lookup(&cast_column, &[10, 20, 30, 100_000]); ++ ++ let predicate = PruningPredicateBuilder::new() ++ .with_file_schema(Arc::clone(&schema)) ++ .try_build(lookup) ++ .unwrap(); ++ ++ let to_file_type = |values: Vec<Option<i64>>| -> ArrayRef { ++ let values: ArrayRef = Arc::new(values.into_iter().collect::<Int64Array>()); ++ arrow::compute::cast(&values, &file_type).unwrap() ++ }; ++ let statistics = TestStatistics::new().with( ++ "b", ++ ContainerStats::new() ++ .with_min(to_file_type(vec![Some(5000), Some(15), Some(50000)])) ++ .with_max(to_file_type(vec![Some(8000), Some(25), Some(60000)])), ++ ); ++ ++ let result = predicate.prune(&statistics).unwrap(); ++ // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 ++ // holds 20. A working cast excludes the first and third; a declined one ++ // falls back to keeping every container. ++ let expected = if pruned { ++ vec![false, true, false] ++ } else { ++ vec![true, true, true] ++ }; ++ assert_eq!(result, expected); ++ } ++ ++ #[test] ++ fn test_partition_routed_hash_lookup_pruning() { ++ let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); ++ let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); ++ let literal = |value: ScalarValue| -> Arc<dyn PhysicalExpr> { ++ Arc::new(phys_expr::Literal::new(value)) ++ }; ++ // One branch per build partition, each holding its own slice of the build ++ // side, as `build_partitioned_filter` produces them. ++ let case: Arc<dyn PhysicalExpr> = Arc::new( ++ phys_expr::CaseExpr::try_new( ++ Some(literal(ScalarValue::UInt64(Some(0)))), ++ vec![ ++ ( ++ literal(ScalarValue::UInt64(Some(0))), ++ hash_lookup(&column, &[10, 20, 30]), ++ ), ++ ( ++ literal(ScalarValue::UInt64(Some(1))), ++ hash_lookup(&column, &[40, 50, 60]), ++ ), ++ ], ++ // Partitions with no build rows reject everything routed to them. ++ Some(literal(ScalarValue::Boolean(Some(false)))), ++ ) ++ .unwrap(), ++ ); ++ ++ let predicate = PruningPredicateBuilder::new() ++ .with_file_schema(Arc::clone(&schema)) ++ .try_build(case) ++ .unwrap(); ++ ++ let statistics = TestStatistics::new().with( ++ "b", ++ ContainerStats::new_i32( ++ vec![Some(5), Some(15), Some(45), Some(100)], ++ vec![Some(8), Some(25), Some(55), Some(200)], ++ ), ++ ); ++ ++ let result = predicate.prune(&statistics).unwrap(); ++ // Container 1 ([15,25]) holds 20 from the first branch and container 2 ++ // ([45,55]) holds 50 from the second, so a branch may match in each. The ++ // other two intersect neither branch, nor the `ELSE`. ++ assert_eq!(result, vec![false, true, true, false]); ++ } + }
diff --git a/native/vendor/datafusion-physical-plan/Cargo.toml b/native/vendor/datafusion-physical-plan/Cargo.toml new file mode 100644 index 0000000..6cf15f7 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/Cargo.toml
@@ -0,0 +1,308 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2024" +rust-version = "1.94.0" +name = "datafusion-physical-plan" +version = "55.1.0" +authors = ["Apache DataFusion <dev@datafusion.apache.org>"] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "Physical (ExecutionPlan) implementations for DataFusion query engine" +homepage = "https://datafusion.apache.org" +readme = "README.md" +keywords = [ + "arrow", + "query", + "sql", +] +license = "Apache-2.0" +repository = "https://github.com/apache/datafusion" +resolver = "2" + +[package.metadata.docs.rs] +all-features = true + +[features] +force_hash_collisions = [] +proto = [ + "dep:datafusion-proto-models", + "dep:datafusion-proto-common", + "datafusion-physical-expr/proto", + "datafusion-physical-expr-common/proto", +] +test_utils = ["arrow/test_utils"] +tokio_coop = [] +tokio_coop_fallback = [] + +[lib] +name = "datafusion_physical_plan" +path = "src/lib.rs" + +[[bench]] +name = "aggregate_vectorized" +path = "benches/aggregate_vectorized.rs" +harness = false +required-features = ["test_utils"] + +[[bench]] +name = "bounded_window" +path = "benches/bounded_window.rs" +harness = false + +[[bench]] +name = "compute_statistics" +path = "benches/compute_statistics.rs" +harness = false + +[[bench]] +name = "dictionary_group_values" +path = "benches/dictionary_group_values.rs" +harness = false + +[[bench]] +name = "hash_join_semi_anti" +path = "benches/hash_join_semi_anti.rs" +harness = false +required-features = ["test_utils"] + +[[bench]] +name = "multi_group_by" +path = "benches/multi_group_by.rs" +harness = false +required-features = ["test_utils"] + +[[bench]] +name = "partial_ordering" +path = "benches/partial_ordering.rs" +harness = false + +[[bench]] +name = "sort_merge_join" +path = "benches/sort_merge_join.rs" +harness = false +required-features = ["test_utils"] + +[[bench]] +name = "sort_preserving_merge" +path = "benches/sort_preserving_merge.rs" +harness = false + +[[bench]] +name = "spill_io" +path = "benches/spill_io.rs" +harness = false + +[dependencies.arrow] +version = "59.2.0" +features = [ + "prettyprint", + "chrono-tz", +] + +[dependencies.arrow-data] +version = "59.2.0" +default-features = false + +[dependencies.arrow-ipc] +version = "59.2.0" +features = [ + "lz4", + "zstd", + "lz4", + "zstd", +] +default-features = false + +[dependencies.arrow-ord] +version = "59.2.0" +default-features = false + +[dependencies.arrow-schema] +version = "59.2.0" +default-features = false + +[dependencies.async-trait] +version = "0.1.89" + +[dependencies.bytes] +version = "1.11" + +[dependencies.datafusion-common] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-common-runtime] +version = "55.1.0" + +[dependencies.datafusion-execution] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-expr] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-functions] +version = "55.1.0" + +[dependencies.datafusion-functions-aggregate-common] +version = "55.1.0" + +[dependencies.datafusion-functions-window-common] +version = "55.1.0" + +[dependencies.datafusion-physical-expr] +version = "55.1.0" +default-features = true + +[dependencies.datafusion-physical-expr-common] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-proto-common] +version = "55.1.0" +optional = true + +[dependencies.datafusion-proto-models] +version = "55.1.0" +optional = true + +[dependencies.futures] +version = "0.3" + +[dependencies.half] +version = "2.7.0" +default-features = false + +[dependencies.hashbrown] +version = "0.17.1" + +[dependencies.indexmap] +version = "2.14.0" + +[dependencies.itertools] +version = "0.15" +features = ["use_std"] + +[dependencies.log] +version = "^0.4" + +[dependencies.num-traits] +version = "0.2" + +[dependencies.parking_lot] +version = "0.12" + +[dependencies.pin-project-lite] +version = "^0.2.7" + +[dependencies.serde_json] +version = "1" +features = ["preserve_order"] + +[dependencies.tokio] +version = "1.52" +features = [ + "macros", + "rt", + "sync", +] + +[dev-dependencies.arrow-data] +version = "59.2.0" +default-features = false + +[dev-dependencies.criterion] +version = "0.8" +features = ["async_futures"] + +[dev-dependencies.datafusion-functions-aggregate] +version = "55.1.0" + +[dev-dependencies.datafusion-functions-window] +version = "55.1.0" + +[dev-dependencies.insta] +version = "1.47.2" +features = [ + "glob", + "filters", +] + +[dev-dependencies.rand] +version = "0.9" + +[dev-dependencies.rstest] +version = "0.26.1" + +[dev-dependencies.rstest_reuse] +version = "0.7.0" + +[dev-dependencies.tokio] +version = "1.52" +features = [ + "macros", + "rt", + "sync", + "rt-multi-thread", + "fs", + "parking_lot", +] + +[lints.clippy] +allow_attributes = "warn" +assigning_clones = "warn" +inefficient_to_string = "warn" +large_futures = "warn" +needless_pass_by_value = "warn" +or_fun_call = "warn" +uninlined_format_args = "warn" +unnecessary_lazy_evaluations = "warn" +unused_async = "warn" +used_underscore_binding = "warn" + +[lints.rust] +# Rust 1.99 renamed AtomicUsize::fetch_update. Retain the upstream implementation +# and its Rust 1.94 minimum instead of requiring the new try_update method. +deprecated = "allow" +unused_qualifications = "deny" + +[lints.rust.unexpected_cfgs] +level = "warn" +priority = 0 +check-cfg = [ + 'cfg(datafusion_coop, values("tokio", "tokio_fallback", "per_stream"))', + "cfg(coverage)", + "cfg(coverage_nightly)", +]
diff --git a/native/vendor/datafusion-physical-plan/Cargo.toml.orig b/native/vendor/datafusion-physical-plan/Cargo.toml.orig new file mode 100644 index 0000000..0f72b74 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/Cargo.toml.orig
@@ -0,0 +1,149 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-physical-plan" +description = "Physical (ExecutionPlan) implementations for DataFusion query engine" +keywords = ["arrow", "query", "sql"] +readme = "README.md" +version = { workspace = true } +edition = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +authors = { workspace = true } +rust-version = { workspace = true } + +[package.metadata.docs.rs] +all-features = true + +# Note: add additional linter rules in lib.rs. +# Rust does not support workspace + new linter rules in subcrates yet +# https://github.com/rust-lang/cargo/issues/13157 +[lints] +workspace = true + +[features] +force_hash_collisions = [] +test_utils = ["arrow/test_utils"] +tokio_coop = [] +tokio_coop_fallback = [] +# Enables `PhysicalExpr::try_to_proto` / `try_from_proto` hooks on the +# physical expressions defined in this crate (e.g. `HashExpr`). Off by +# default so consumers that never serialize plans pay nothing. +proto = [ + "dep:datafusion-proto-models", + "dep:datafusion-proto-common", + "datafusion-physical-expr/proto", + "datafusion-physical-expr-common/proto", +] + +[lib] +name = "datafusion_physical_plan" + +[dependencies] +arrow = { workspace = true } +arrow-data = { workspace = true } +# Spill IPC writes require lz4 and zstd codec support. Keep these features in +# sync with the SpillCompression variants in datafusion-common so codec +# availability is explicit in the crate that owns spill handling. +arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } +arrow-ord = { workspace = true } +arrow-schema = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +datafusion-common = { workspace = true } +datafusion-common-runtime = { workspace = true, default-features = true } +datafusion-execution = { workspace = true } +datafusion-expr = { workspace = true } +datafusion-functions = { workspace = true } +datafusion-functions-aggregate-common = { workspace = true } +datafusion-functions-window-common = { workspace = true } +datafusion-physical-expr = { workspace = true, default-features = true } +datafusion-physical-expr-common = { workspace = true } +datafusion-proto-common = { workspace = true, optional = true } +datafusion-proto-models = { workspace = true, optional = true } +futures = { workspace = true } +half = { workspace = true } +hashbrown = { workspace = true } +indexmap = { workspace = true } +itertools = { workspace = true, features = ["use_std"] } +log = { workspace = true } +num-traits = { workspace = true } +parking_lot = { workspace = true } +pin-project-lite = { workspace = true } +serde_json = { workspace = true, features = ["preserve_order"] } +tokio = { workspace = true } + +[dev-dependencies] +arrow-data = { workspace = true } +criterion = { workspace = true, features = ["async_futures"] } +datafusion-functions-aggregate = { workspace = true } +datafusion-functions-window = { workspace = true } +insta = { workspace = true } +rand = { workspace = true } +rstest = { workspace = true } +rstest_reuse = "0.7.0" +tokio = { workspace = true, features = [ + "rt-multi-thread", + "fs", + "parking_lot", +] } + +[[bench]] +harness = false +name = "partial_ordering" + +[[bench]] +harness = false +name = "spill_io" + +[[bench]] +harness = false +name = "sort_preserving_merge" + +[[bench]] +harness = false +name = "sort_merge_join" +required-features = ["test_utils"] + +[[bench]] +harness = false +name = "aggregate_vectorized" +required-features = ["test_utils"] + +[[bench]] +harness = false +name = "compute_statistics" + +[[bench]] +harness = false +name = "dictionary_group_values" + +[[bench]] +harness = false +name = "hash_join_semi_anti" +required-features = ["test_utils"] + +[[bench]] +harness = false +name = "multi_group_by" +required-features = ["test_utils"] + +[[bench]] +harness = false +name = "bounded_window"
diff --git a/native/vendor/datafusion-physical-plan/LICENSE.txt b/native/vendor/datafusion-physical-plan/LICENSE.txt new file mode 100644 index 0000000..d74c6b5 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/LICENSE.txt
@@ -0,0 +1,212 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + +This project includes code from Apache Aurora. + +* dev/release/{release,changelog,release-candidate} are based on the scripts from + Apache Aurora + +Copyright: 2016 The Apache Software Foundation. +Home page: https://aurora.apache.org/ +License: http://www.apache.org/licenses/LICENSE-2.0
diff --git a/native/vendor/datafusion-physical-plan/NOTICE.txt b/native/vendor/datafusion-physical-plan/NOTICE.txt new file mode 100644 index 0000000..0bd2d52 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/NOTICE.txt
@@ -0,0 +1,5 @@ +Apache DataFusion +Copyright 2019-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/).
diff --git a/native/vendor/datafusion-physical-plan/README.md b/native/vendor/datafusion-physical-plan/README.md new file mode 100644 index 0000000..3a33100 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/README.md
@@ -0,0 +1,33 @@ +<!--- + 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. +--> + +# Apache DataFusion Physical Plan + +[Apache DataFusion] is an extensible query execution framework, written in Rust, that uses [Apache Arrow] as its in-memory format. + +This crate is a submodule of DataFusion that contains the `ExecutionPlan` trait and the various implementations of that +trait for built in operators such as filters, projections, joins, aggregations, etc. + +Most projects should use the [`datafusion`] crate directly, which re-exports +this module. If you are already using the [`datafusion`] crate, there is no +reason to use this crate directly in your project as well. + +[apache arrow]: https://arrow.apache.org/ +[apache datafusion]: https://datafusion.apache.org/ +[`datafusion`]: https://crates.io/crates/datafusion
diff --git a/native/vendor/datafusion-physical-plan/benches/aggregate_vectorized.rs b/native/vendor/datafusion-physical-plan/benches/aggregate_vectorized.rs new file mode 100644 index 0000000..488647d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/aggregate_vectorized.rs
@@ -0,0 +1,309 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, BooleanBufferBuilder}; +use arrow::datatypes::{Int32Type, StringViewType}; +use arrow::util::bench_util::{ + create_primitive_array, create_string_view_array_with_len, + create_string_view_array_with_max_len, +}; +use arrow_schema::DataType; +use criterion::measurement::WallTime; +use criterion::{ + BenchmarkGroup, BenchmarkId, Criterion, criterion_group, criterion_main, +}; +use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupColumn; +use datafusion_physical_plan::aggregates::group_values::multi_group_by::bytes_view::ByteViewGroupValueBuilder; +use datafusion_physical_plan::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; +use rand::SeedableRng; +use rand::distr::{Bernoulli, Distribution}; +use rand::rngs::StdRng; +use std::hint::black_box; +use std::sync::Arc; + +const SIZES: [usize; 3] = [1_000, 10_000, 100_000]; +const NULL_DENSITIES: [f32; 3] = [0.0, 0.1, 0.5]; + +fn bench_vectorized_append(c: &mut Criterion) { + byte_view_vectorized_append(c); + primitive_vectorized_append(c); +} + +fn byte_view_vectorized_append(c: &mut Criterion) { + let mut group = c.benchmark_group("ByteViewGroupValueBuilder_vectorized_append"); + + for &size in &SIZES { + let rows: Vec<usize> = (0..size).collect(); + + for &null_density in &NULL_DENSITIES { + let input = create_string_view_array_with_len(size, null_density, 8, false); + let input: ArrayRef = Arc::new(input); + + bytes_bench(&mut group, "inline", size, &rows, null_density, &input); + } + } + + for &size in &SIZES { + let rows: Vec<usize> = (0..size).collect(); + + for &null_density in &NULL_DENSITIES { + let input = create_string_view_array_with_len(size, null_density, 64, true); + let input: ArrayRef = Arc::new(input); + + bytes_bench(&mut group, "scenario", size, &rows, null_density, &input); + } + } + + for &size in &SIZES { + let rows: Vec<usize> = (0..size).collect(); + + for &null_density in &NULL_DENSITIES { + let input = create_string_view_array_with_max_len(size, null_density, 400); + let input: ArrayRef = Arc::new(input); + + bytes_bench(&mut group, "random", size, &rows, null_density, &input); + } + } + + group.finish(); +} + +fn bytes_bench( + group: &mut BenchmarkGroup<WallTime>, + bench_prefix: &str, + size: usize, + rows: &Vec<usize>, + null_density: f32, + input: &ArrayRef, +) { + // vectorized_append + let function_name = format!("{bench_prefix}_null_{null_density:.1}_size_{size}"); + let id = BenchmarkId::new(&function_name, "vectorized_append"); + group.bench_function(id, |b| { + b.iter(|| { + let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new(); + builder.vectorized_append(input, rows).unwrap(); + }); + }); + + // append_val + let id = BenchmarkId::new(&function_name, "append_val"); + group.bench_function(id, |b| { + b.iter(|| { + let mut builder = ByteViewGroupValueBuilder::<StringViewType>::new(); + for &i in rows { + builder.append_val(input, i).unwrap(); + } + }); + }); + + // vectorized_equal_to + vectorized_equal_to( + group, + ByteViewGroupValueBuilder::<StringViewType>::new(), + &function_name, + rows, + input, + "all_true", + vec![true; size], + ); + vectorized_equal_to( + group, + ByteViewGroupValueBuilder::<StringViewType>::new(), + &function_name, + rows, + input, + "0.75 true", + { + let mut rng = StdRng::seed_from_u64(42); + let d = Bernoulli::new(0.75).unwrap(); + (0..size).map(|_| d.sample(&mut rng)).collect::<Vec<_>>() + }, + ); + vectorized_equal_to( + group, + ByteViewGroupValueBuilder::<StringViewType>::new(), + &function_name, + rows, + input, + "0.5 true", + { + let mut rng = StdRng::seed_from_u64(42); + let d = Bernoulli::new(0.5).unwrap(); + (0..size).map(|_| d.sample(&mut rng)).collect::<Vec<_>>() + }, + ); + vectorized_equal_to( + group, + ByteViewGroupValueBuilder::<StringViewType>::new(), + &function_name, + rows, + input, + "0.25 true", + { + let mut rng = StdRng::seed_from_u64(42); + let d = Bernoulli::new(0.25).unwrap(); + (0..size).map(|_| d.sample(&mut rng)).collect::<Vec<_>>() + }, + ); + // Not adding 0 true case here as if we optimize for 0 true cases the caller should avoid calling this method at all +} + +fn primitive_vectorized_append(c: &mut Criterion) { + let mut group = c.benchmark_group("PrimitiveGroupValueBuilder_vectorized_append"); + + for &size in &SIZES { + let rows: Vec<usize> = (0..size).collect(); + + for &null_density in &NULL_DENSITIES { + if null_density == 0.0 { + bench_single_primitive::<false>(&mut group, size, &rows, null_density) + } + bench_single_primitive::<true>(&mut group, size, &rows, null_density); + } + } + + group.finish(); +} + +fn bench_single_primitive<const NULLABLE: bool>( + group: &mut BenchmarkGroup<WallTime>, + size: usize, + rows: &Vec<usize>, + null_density: f32, +) { + if !NULLABLE { + assert_eq!( + null_density, 0.0, + "non-nullable case must have null_density 0" + ); + } + + let input = create_primitive_array::<Int32Type>(size, null_density); + let input: ArrayRef = Arc::new(input); + let function_name = format!("null_{null_density:.1}_nullable_{NULLABLE}_size_{size}"); + + // vectorized_append + let id = BenchmarkId::new(&function_name, "vectorized_append"); + group.bench_function(id, |b| { + b.iter(|| { + let mut builder = + PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32); + builder.vectorized_append(&input, rows).unwrap(); + }); + }); + + // append_val + let id = BenchmarkId::new(&function_name, "append_val"); + group.bench_function(id, |b| { + b.iter(|| { + let mut builder = + PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32); + for &i in rows { + builder.append_val(&input, i).unwrap(); + } + }); + }); + + // vectorized_equal_to + vectorized_equal_to( + group, + PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32), + &function_name, + rows, + &input, + "all_true", + vec![true; size], + ); + vectorized_equal_to( + group, + PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32), + &function_name, + rows, + &input, + "0.75 true", + { + let mut rng = StdRng::seed_from_u64(42); + let d = Bernoulli::new(0.75).unwrap(); + (0..size).map(|_| d.sample(&mut rng)).collect::<Vec<_>>() + }, + ); + vectorized_equal_to( + group, + PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32), + &function_name, + rows, + &input, + "0.5 true", + { + let mut rng = StdRng::seed_from_u64(42); + let d = Bernoulli::new(0.5).unwrap(); + (0..size).map(|_| d.sample(&mut rng)).collect::<Vec<_>>() + }, + ); + vectorized_equal_to( + group, + PrimitiveGroupValueBuilder::<Int32Type, NULLABLE>::new(DataType::Int32), + &function_name, + rows, + &input, + "0.25 true", + { + let mut rng = StdRng::seed_from_u64(42); + let d = Bernoulli::new(0.25).unwrap(); + (0..size).map(|_| d.sample(&mut rng)).collect::<Vec<_>>() + }, + ); + // Not adding 0 true case here as if we optimize for 0 true cases the caller should avoid calling this method at all +} + +/// Test `vectorized_equal_to` with different number of true in the initial results +#[expect(clippy::needless_pass_by_value)] +fn vectorized_equal_to<GroupColumnBuilder: GroupColumn>( + group: &mut BenchmarkGroup<WallTime>, + mut builder: GroupColumnBuilder, + function_name: &str, + rows: &[usize], + input: &ArrayRef, + equal_to_result_description: &str, + equal_to_results: Vec<bool>, +) { + let id = BenchmarkId::new( + function_name, + format!("vectorized_equal_to_{equal_to_result_description}"), + ); + group.bench_function(id, |b| { + builder.vectorized_append(input, rows).unwrap(); + + b.iter(|| { + // Rebuild the buffer each iteration as `vectorized_equal_to` mutates + // it, and without a fresh buffer all iterations after the first one + // would not be meaningful. + let mut equal_to_buffer = BooleanBufferBuilder::new(equal_to_results.len()); + for &v in &equal_to_results { + equal_to_buffer.append(v); + } + builder.vectorized_equal_to(rows, input, rows, &mut equal_to_buffer); + + // Make sure that the compiler does not optimize away the call + black_box(equal_to_buffer); + }); + }); +} + +criterion_group!(benches, bench_vectorized_append); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/bounded_window.rs b/native/vendor/datafusion-physical-plan/benches/bounded_window.rs new file mode 100644 index 0000000..56e195a --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/bounded_window.rs
@@ -0,0 +1,280 @@ +// 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 `BoundedWindowAggExec` with many partitions. +//! +//! The streaming window operator keeps per-partition state keyed by +//! `PartitionKey` (`Vec<ScalarValue>`) and, in `Linear` mode (input sorted +//! by the ORDER BY column but not by the partition columns), visits every +//! live partition on every batch while never retiring partitions until the +//! input is exhausted. The cases here stress that path in different ways: +//! +//! - `linear N partitions`: dense round-robin keys -- every partition +//! receives rows in every batch, so per-visit fixed costs dominate. +//! - `linear sparse N partitions`: keys are clustered in time, so each +//! batch touches only a small, fresh subset of keys while the set of live +//! partitions keeps growing -- per-batch work on quiet partitions +//! dominates. +//! - `linear rows N partitions`: the dense layout with a ROWS frame, whose +//! results can only be finalized as more rows of the same partition +//! arrive. +//! - `linear multi N partitions`: two window expressions over the dense +//! layout, doubling the per-partition evaluation sweeps. +//! - `sorted N partitions`: control; input sorted by partition key, so +//! finished partitions are pruned eagerly and the state maps stay small. + +use std::sync::Arc; + +use arrow::array::UInt64Array; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_execution::TaskContext; +use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, +}; +use datafusion_functions_aggregate::count::count_udaf; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr}; +use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, collect}; + +const BATCH_SIZE: usize = 8192; +const N_BATCHES: usize = 16; +/// Distinct partition keys per batch in the sparse layout. Each batch +/// introduces this many previously-unseen keys, so the total partition count +/// is `N_BATCHES * SPARSE_KEYS_PER_BATCH`. +const SPARSE_KEYS_PER_BATCH: usize = 2048; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("pk", DataType::UInt64, false), + Field::new("ts", DataType::UInt64, false), + ])) +} + +/// Batches with `ts` ascending across the whole input and partition keys +/// chosen by `pk_of_row`. +fn make_batches(pk_of_row: impl Fn(usize) -> u64) -> Vec<RecordBatch> { + (0..N_BATCHES) + .map(|b| { + let start = b * BATCH_SIZE; + let pk: UInt64Array = (start..start + BATCH_SIZE) + .map(|i| Some(pk_of_row(i))) + .collect(); + let ts: UInt64Array = (start..start + BATCH_SIZE) + .map(|i| Some(i as u64)) + .collect(); + RecordBatch::try_new(schema(), vec![Arc::new(pk), Arc::new(ts)]).unwrap() + }) + .collect() +} + +/// Round-robin over `n_partitions`: every partition receives rows in every +/// batch (when `n_partitions <= BATCH_SIZE`). +fn dense_batches(n_partitions: usize) -> Vec<RecordBatch> { + make_batches(move |i| (i % n_partitions) as u64) +} + +/// Keys clustered in time: batch `b` only contains keys in +/// `[b * SPARSE_KEYS_PER_BATCH, (b + 1) * SPARSE_KEYS_PER_BATCH)`, cycled so +/// that consecutive rows belong to different partitions. Previously-seen +/// keys never recur, but `Linear` mode cannot know that, so the live +/// partition set grows for the whole run. +fn sparse_batches() -> Vec<RecordBatch> { + make_batches(|i| { + ((i / BATCH_SIZE) * SPARSE_KEYS_PER_BATCH + (i % SPARSE_KEYS_PER_BATCH)) as u64 + }) +} + +/// Input laid out partition-by-partition (the `Sorted` layout). +fn sorted_batches(n_partitions: usize) -> Vec<RecordBatch> { + let rows_per_partition = BATCH_SIZE * N_BATCHES / n_partitions; + make_batches(move |i| (i / rows_per_partition) as u64) +} + +fn sort_expr(name: &str) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col(name, &schema()).unwrap(), + options: Default::default(), + } +} + +/// `RANGE BETWEEN CURRENT ROW AND 10 FOLLOWING` +fn range_frame() -> WindowFrame { + WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(Some(10))), + ) +} + +/// `ROWS BETWEEN CURRENT ROW AND 2 FOLLOWING` +fn rows_frame() -> WindowFrame { + WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(Some(2))), + ) +} + +/// `<agg>(ts) OVER (PARTITION BY pk ORDER BY ts <window_frame>)` for each +/// aggregate in `aggregates`. +fn window_exec( + batches: Vec<RecordBatch>, + mode: InputOrderMode, + input_ordering: Vec<PhysicalSortExpr>, + window_frame: &WindowFrame, + aggregates: &[(WindowFunctionDefinition, &str)], +) -> Arc<dyn ExecutionPlan> { + let schema = schema(); + let source = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None) + .expect("memory exec") + .try_with_sort_information(LexOrdering::new(input_ordering).into_iter().collect()) + .expect("sort information"); + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(source))); + let args = vec![col("ts", &schema).unwrap()]; + let partitionby_exprs = vec![col("pk", &schema).unwrap()]; + let orderby_exprs = vec![PhysicalSortExpr { + expr: col("ts", &schema).unwrap(), + options: Default::default(), + }]; + let window_expr = aggregates + .iter() + .map(|(fun, name)| { + create_window_expr( + fun, + name.to_string(), + &args, + &partitionby_exprs, + &orderby_exprs, + Arc::new(window_frame.clone()), + input.schema(), + false, + false, + None, + ) + .expect("window expr") + }) + .collect::<Vec<_>>(); + Arc::new( + BoundedWindowAggExec::try_new(window_expr, input, mode, true) + .expect("bounded window exec"), + ) +} + +fn count() -> (WindowFunctionDefinition, &'static str) { + ( + WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count", + ) +} + +fn sum() -> (WindowFunctionDefinition, &'static str) { + (WindowFunctionDefinition::AggregateUDF(sum_udaf()), "sum") +} + +fn bounded_window_benchmark(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bounded_window_partitions"); + group.sample_size(10); + + let mut run_case = |name: String, plan: Arc<dyn ExecutionPlan>| { + group.bench_function(name, |b| { + b.iter(|| { + let task_ctx = Arc::new(TaskContext::default()); + let batches = rt + .block_on(collect(Arc::clone(&plan), task_ctx)) + .expect("execution"); + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::<usize>(), + BATCH_SIZE * N_BATCHES + ); + }) + }); + }; + + for n_partitions in [100, 10_000] { + run_case( + format!("linear {n_partitions} partitions"), + window_exec( + dense_batches(n_partitions), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &range_frame(), + &[count()], + ), + ); + } + + run_case( + format!( + "linear sparse {} partitions", + N_BATCHES * SPARSE_KEYS_PER_BATCH + ), + window_exec( + sparse_batches(), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &range_frame(), + &[count()], + ), + ); + + run_case( + "linear rows 10000 partitions".to_string(), + window_exec( + dense_batches(10_000), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &rows_frame(), + &[count()], + ), + ); + + run_case( + "linear multi 10000 partitions".to_string(), + window_exec( + dense_batches(10_000), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &range_frame(), + &[count(), sum()], + ), + ); + + // Control: the same query over partition-sorted input, where finished + // partitions are pruned eagerly and the state maps stay small. + run_case( + "sorted 10000 partitions".to_string(), + window_exec( + sorted_batches(10_000), + InputOrderMode::Sorted, + vec![sort_expr("pk"), sort_expr("ts")], + &range_frame(), + &[count()], + ), + ); + + group.finish(); +} + +criterion_group!(benches, bounded_window_benchmark); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/compute_statistics.rs b/native/vendor/datafusion-physical-plan/benches/compute_statistics.rs new file mode 100644 index 0000000..cddf4c2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/compute_statistics.rs
@@ -0,0 +1,354 @@ +// 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 `compute_statistics` with `StatsCache`. +//! +//! Demonstrates that caching eliminates redundant subtree walks in plans +//! containing partition-merging operators (CoalescePartitionsExec) and +//! binary join trees (CrossJoinExec). +//! +//! The plan shapes here mirror the reproducers from the planning-speed +//! EPIC (<https://github.com/apache/datafusion/issues/19795>): +//! - Coalesce chain: deep linear plans (e.g. deeply nested subqueries) +//! - Cross-join tree: balanced binary trees from multi-way joins +//! (mirrors the `physical_many_self_joins` sql_planner benchmark) + +use std::fmt; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, Statistics}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::Literal; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, ExecutionPlan, PlanProperties, +}; +use datafusion_physical_plan::filter::FilterExec; +use datafusion_physical_plan::joins::CrossJoinExec; +use datafusion_physical_plan::statistics::StatisticsArgs; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Partitioning, + ReplaceChildrenOptions, SendableRecordBatchStream, StatisticsContext, +}; + +/// Minimal leaf node for benchmarking +#[derive(Debug)] +struct BenchLeaf { + schema: SchemaRef, + cache: Arc<PlanProperties>, +} + +impl BenchLeaf { + fn new(col_name: &str) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + col_name, + DataType::Int32, + false, + )])); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(2), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { schema, cache } + } +} + +impl DisplayAs for BenchLeaf { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "BenchLeaf") + } +} + +impl ExecutionPlan for BenchLeaf { + fn name(&self) -> &str { + "BenchLeaf" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } +} + +/// Build: CoalescePartitions^depth -> BenchLeaf +fn build_coalesce_chain(depth: usize) -> Arc<dyn ExecutionPlan> { + let mut plan: Arc<dyn ExecutionPlan> = Arc::new(BenchLeaf::new("a")); + for _ in 0..depth { + plan = Arc::new(CoalescePartitionsExec::new(plan)); + } + plan +} + +/// Build a balanced binary tree of CrossJoinExec with 2^depth leaves. +/// Mirrors the plan shape produced by multi-way self-joins like the +/// `physical_many_self_joins` benchmark in sql_planner.rs (#19795). +fn build_cross_join_tree(depth: usize, next_col: &mut usize) -> Arc<dyn ExecutionPlan> { + if depth == 0 { + let col_name = format!("c{next_col}"); + *next_col += 1; + return Arc::new(BenchLeaf::new(&col_name)); + } + let left = build_cross_join_tree(depth - 1, next_col); + let right = build_cross_join_tree(depth - 1, next_col); + Arc::new(CrossJoinExec::new(left, right)) +} + +/// Build: Filter^depth -> BenchLeaf (always-true predicate). +fn build_filter_chain(depth: usize) -> Arc<dyn ExecutionPlan> { + let mut plan: Arc<dyn ExecutionPlan> = Arc::new(BenchLeaf::new("a")); + let predicate: Arc<dyn PhysicalExpr> = + Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))); + for _ in 0..depth { + plan = Arc::new( + FilterExec::try_new(Arc::clone(&predicate), plan) + .expect("FilterExec::try_new failed"), + ); + } + plan +} + +/// Build a mixed chain alternating partition-merging and partition-preserving +/// operators: (Coalesce -> Filter -> Filter) repeated `groups` times -> BenchLeaf. +/// Exercises the cache with both None and Some(p) lookups in the same walk. +fn build_mixed_chain(groups: usize) -> Arc<dyn ExecutionPlan> { + let mut plan: Arc<dyn ExecutionPlan> = Arc::new(BenchLeaf::new("a")); + let predicate: Arc<dyn PhysicalExpr> = + Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))); + for _ in 0..groups { + // Two partition-preserving filters + for _ in 0..2 { + plan = Arc::new( + FilterExec::try_new(Arc::clone(&predicate), plan) + .expect("FilterExec::try_new failed"), + ); + } + // One partition-merging coalesce + plan = Arc::new(CoalescePartitionsExec::new(plan)); + } + plan +} + +/// Recursive walk without a shared cross-node cache, simulating pre-cache behavior. +/// Each node is computed with a fresh `StatisticsContext`, so every call triggers a +/// fresh subtree walk, resulting in O(n^2) total node visits for a chain of depth n. +/// +/// Note: each `StatisticsContext::compute` re-walk still benefits from its own +/// ephemeral cache; only the cross-node sharing is removed. +fn compute_statistics_without_shared_cache( + plan: &dyn ExecutionPlan, + partition: Option<usize>, +) -> Result<Arc<Statistics>> { + for child in plan.children() { + compute_statistics_without_shared_cache(child.as_ref(), None)?; + } + let args = StatisticsArgs::new().with_partition(partition); + StatisticsContext::new().compute(plan, &args) +} + +fn bench_compute_statistics(c: &mut Criterion) { + // --- Coalesce chain (linear plan) --- + // Deep linear plans arise from deeply nested subqueries, CTEs, etc. + let mut group = c.benchmark_group("compute_statistics_coalesce_chain"); + for depth in [10, 20, 50] { + let plan = build_coalesce_chain(depth); + group.bench_with_input(BenchmarkId::new("cached", depth), &plan, |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute(plan.as_ref(), &StatisticsArgs::new()) + .unwrap() + }); + }); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", depth), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), None).unwrap() + }); + }, + ); + } + group.finish(); + + // --- Cross-join tree (balanced binary plan) --- + // Binary trees arise from multi-way joins (e.g. physical_many_self_joins + // in sql_planner.rs, see #19795). CrossJoinExec calls + // StatisticsContext::compute for per-partition stats, re-walking the left + // subtree at each node. The gap between cached/uncached is smaller than + // the linear chain because only the left child triggers a re-walk. + let mut group = c.benchmark_group("compute_statistics_cross_join_tree"); + for depth in [3, 5, 7] { + let mut next_col = 0; + let plan = build_cross_join_tree(depth, &mut next_col); + let label = format!("depth={depth}_leaves={}", 1usize << depth); + group.bench_with_input(BenchmarkId::new("cached", &label), &plan, |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap() + }); + }); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", &label), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), Some(0)) + .unwrap() + }); + }, + ); + } + group.finish(); + + // --- Filter chain (partition-preserving linear plan) --- + // When called with Some(0), the framework first walks the entire tree + // computing None stats, then each filter requests Some(0) on demand. + // Both walks are cached, so the total cost is ~2n vs n node visits for None. + let mut group = c.benchmark_group("compute_statistics_filter_chain"); + for depth in [10, 20, 50] { + let plan = build_filter_chain(depth); + group.bench_with_input( + BenchmarkId::new("cached_partition", depth), + &plan, + |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap() + }); + }, + ); + group.bench_with_input( + BenchmarkId::new("cached_overall", depth), + &plan, + |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute(plan.as_ref(), &StatisticsArgs::new()) + .unwrap() + }); + }, + ); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", depth), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), Some(0)) + .unwrap() + }); + }, + ); + } + group.finish(); + + // --- Mixed chain (partition-preserving + partition-merging) --- + // Alternates Filter (preserving) and CoalescePartitions (merging) to + // exercise the cache with both None and Some(p) lookups in a single walk. + let mut group = c.benchmark_group("compute_statistics_mixed_chain"); + for groups in [3, 5, 10] { + let plan = build_mixed_chain(groups); + let depth = groups * 3; // 2 filters + 1 coalesce per group + group.bench_with_input(BenchmarkId::new("cached", depth), &plan, |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap() + }); + }); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", depth), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), Some(0)) + .unwrap() + }); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_compute_statistics); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/dictionary_group_values.rs b/native/vendor/datafusion-physical-plan/benches/dictionary_group_values.rs new file mode 100644 index 0000000..ded52ae --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/dictionary_group_values.rs
@@ -0,0 +1,176 @@ +// 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 `GroupValues` over a single `Dictionary<Int32, Utf8>` +//! column. Each iteration measures `intern` (once or N times) followed by +//! `emit(EmitTo::All)`. The `Box<dyn GroupValues>` returned by +//! `new_group_values` is constructed in the setup closure of +//! `iter_batched_ref` and is not included in the timing. + +use arrow::array::{ArrayRef, DictionaryArray, PrimitiveArray, StringArray}; +use arrow::buffer::{Buffer, NullBuffer}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema, SchemaRef}; +use criterion::{ + BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main, +}; +use datafusion_expr::EmitTo; +use datafusion_physical_plan::aggregates::group_values::new_group_values; +use datafusion_physical_plan::aggregates::order::GroupOrdering; +use rand::rngs::StdRng; +use rand::seq::SliceRandom; +use rand::{Rng, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +const SIZES: [usize; 2] = [8 * 1024, 64 * 1024]; +const CARDS_RELATIVE: [usize; 4] = [20, 75, 300, 1000]; +const N_BATCHES: usize = 4; +// Fixed for reproducibility. +const SEED: u64 = 0xD1C7; + +fn dict_schema() -> SchemaRef { + let dict_ty = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + Arc::new(Schema::new(vec![Field::new("g", dict_ty, true)])) +} + +/// Build a `Dictionary<Int32, Utf8>` column. +fn make_dict(size: usize, cardinality: usize, null_density: f32, seed: u64) -> ArrayRef { + let strings: Vec<String> = (0..cardinality).map(|i| format!("v_{i:08}")).collect(); + let values = Arc::new(StringArray::from( + strings.iter().map(String::as_str).collect::<Vec<_>>(), + )); + + let mut rng = StdRng::seed_from_u64(seed); + let keys: Vec<i32> = if cardinality == size { + let mut perm: Vec<i32> = (0..size as i32).collect(); + perm.shuffle(&mut rng); + perm + } else { + (0..size) + .map(|_| rng.random_range(0..cardinality) as i32) + .collect() + }; + let keys_buf = Buffer::from_slice_ref(&keys); + + let nulls: Option<NullBuffer> = (null_density > 0.0).then(|| { + (0..size) + .map(|_| !rng.random_bool(null_density as f64)) + .collect() + }); + + let key_array = PrimitiveArray::<Int32Type>::new(keys_buf.into(), nulls); + Arc::new(DictionaryArray::<Int32Type>::try_new(key_array, values).unwrap()) +} + +fn bench_id( + label: &str, + size: usize, + cardinality: usize, + null_density: f32, +) -> BenchmarkId { + BenchmarkId::new( + label, + format!("size_{size}_card_{cardinality}_null_{null_density:.2}"), + ) +} + +fn bench_intern_emit(c: &mut Criterion) { + let mut group = c.benchmark_group("dict_intern_emit"); + let schema = dict_schema(); + let null_density = 0.0; + + for &size in &SIZES { + let mut cards = CARDS_RELATIVE.to_vec(); + cards.push(size); // all-unique stress case + for cardinality in cards { + let array = make_dict(size, cardinality, null_density, SEED); + group.throughput(Throughput::Elements(size as u64)); + group.bench_function( + bench_id("intern_emit", size, cardinality, null_density), + |b| { + b.iter_batched_ref( + || { + ( + new_group_values(schema.clone(), &GroupOrdering::None) + .unwrap(), + Vec::<usize>::with_capacity(size), + ) + }, + |(gv, groups)| { + gv.intern(std::slice::from_ref(&array), groups).unwrap(); + black_box(&*groups); + black_box(gv.emit(EmitTo::All).unwrap()); + }, + BatchSize::SmallInput, + ); + }, + ); + } + } + group.finish(); +} + +fn bench_repeated_intern_emit(c: &mut Criterion) { + let mut group = c.benchmark_group("dict_repeated_intern_emit"); + let schema = dict_schema(); + let null_density = 0.10; + + for &size in &SIZES { + let mut cards = CARDS_RELATIVE.to_vec(); + cards.push(size); + for cardinality in cards { + let batches: Vec<ArrayRef> = (0..N_BATCHES) + .map(|i| { + make_dict( + size, + cardinality, + null_density, + SEED.wrapping_add(i as u64), + ) + }) + .collect(); + group.throughput(Throughput::Elements((size * N_BATCHES) as u64)); + group.bench_function( + bench_id("repeated_intern_emit", size, cardinality, null_density), + |b| { + b.iter_batched_ref( + || { + ( + new_group_values(schema.clone(), &GroupOrdering::None) + .unwrap(), + Vec::<usize>::with_capacity(size), + ) + }, + |(gv, groups)| { + for arr in &batches { + gv.intern(std::slice::from_ref(arr), groups).unwrap(); + black_box(&*groups); + } + black_box(gv.emit(EmitTo::All).unwrap()); + }, + BatchSize::SmallInput, + ); + }, + ); + } + } + group.finish(); +} + +criterion_group!(benches, bench_intern_emit, bench_repeated_intern_emit); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/hash_join_semi_anti.rs b/native/vendor/datafusion-physical-plan/benches/hash_join_semi_anti.rs new file mode 100644 index 0000000..1e11da3 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/hash_join_semi_anti.rs
@@ -0,0 +1,387 @@ +// 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. + +//! Criterion benchmarks for Hash Join with RightSemi/RightAnti joins with Int32 keys. +//! +//! ## Key Benchmark Axes +//! +//! - **Density**: How tightly distinct keys pack into their numeric range. +//! `density = num_distinct_keys / (max_key - min_key + 1)`. +//! Examples for 5 distinct keys: +//! - `[0, 1, 2, 3, 4]` → 5/5 = 100% (fully packed) +//! - `[0, 2, 4, 6, 8]` → 5/9 ≈ 55% (every 2nd slot) +//! - `[0, 10, 20, 30, 40]` → 5/41 ≈ 12% (every 10th slot) +//! +//! Why it matters for this workload: future potential semi/anti-join +//! fast paths could exploit densely packed build keys to outperform the +//! general hash-table path, which is largely insensitive to density. +//! Varying density across benchmarks helps surface those potential gains +//! under different key distributions. Density describes only the +//! build-side key layout; the per-probe match count is tracked +//! separately as fanout. +//! +//! - **Hit Rate**: The percentage of probe rows that find a match in the build side. +//! This controls how often the join produces output rows. +//! +//! Semi/anti joins can short-circuit after finding the first match, so these +//! benchmarks help evaluate optimization strategies for existence checks. + +use std::sync::Arc; + +use arrow::array::{Int32Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::{JoinType, NullEquality}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::collect; +use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, utils::JoinOn}; +use datafusion_physical_plan::test::TestMemoryExec; +use tokio::runtime::Runtime; + +/// Build RecordBatches with Int32 keys. +/// +/// Schema: (key: Int32, data: Int32, payload: Utf8) +/// +/// `key_mod` controls distinct key count: key = row_index % key_mod. +/// `key_offset` shifts keys to control hit rate. +fn build_batches( + num_rows: usize, + key_mod: usize, + key_offset: i32, + schema: &SchemaRef, +) -> Vec<RecordBatch> { + let keys: Vec<i32> = (0..num_rows) + .map(|i| ((i % key_mod) as i32) + key_offset) + .collect(); + let data: Vec<i32> = (0..num_rows).map(|i| i as i32).collect(); + let payload: Vec<String> = data.iter().map(|d| format!("val_{d}")).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(data)), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn make_exec( + batches: &[RecordBatch], + schema: &SchemaRef, +) -> Arc<dyn datafusion_physical_plan::ExecutionPlan> { + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("data", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + ])) +} + +fn do_hash_join( + left: Arc<dyn datafusion_physical_plan::ExecutionPlan>, + right: Arc<dyn datafusion_physical_plan::ExecutionPlan>, + join_type: JoinType, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + + let task_ctx = Arc::new(TaskContext::default()); + rt.block_on(async { + let batches = collect(Arc::new(join), task_ctx).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }) +} + +/// Build batches with sparse keys (key = row_index % key_mod * multiplier + key_offset). +/// The `multiplier` controls density: 1 = 100%, 2 = 50%, 10 = 10%. +fn build_batches_sparse( + num_rows: usize, + key_mod: usize, + key_offset: i32, + multiplier: i32, + schema: &SchemaRef, +) -> Vec<RecordBatch> { + let keys: Vec<i32> = (0..num_rows) + .map(|i| ((i % key_mod) as i32) * multiplier + key_offset) + .collect(); + let data: Vec<i32> = (0..num_rows).map(|i| i as i32).collect(); + let payload: Vec<String> = data.iter().map(|d| format!("val_{d}")).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(data)), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn bench_hash_join_semi_anti(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + + let mut group = c.benchmark_group("hash_join_semi_anti"); + + // Build side: 100K rows, Probe side: 1M rows + // Matching ratio: 1:1 (build keys are unique, each probe matches at most 1 build row) + let build_rows = 100_000; + let probe_rows = 1_000_000; + + // ========================================================================= + // RightSemi Join benchmarks + // ========================================================================= + + // RightSemi - 100% Density, 100% hit rate + // Keys: 0..100K contiguous, all probe rows find a match + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows, 0, &s); + group.bench_function(BenchmarkId::new("right_semi_d100_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 100% Density, 10% hit rate + // Keys: 0..100K contiguous, only 10% of probe rows find a match + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows * 10, 0, &s); + group.bench_function(BenchmarkId::new("right_semi_d100_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 50% Density, 100% hit rate + // Keys: 0, 2, 4, ... (sparse, multiplier=2), all probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_semi_d50_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 50% Density, 10% hit rate + // Keys: 0, 2, 4, ... (sparse), only 10% of probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_semi_d50_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 10% Density, 100% hit rate + // Keys: 0, 10, 20, ... (very sparse, multiplier=10), all probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_semi_d10_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 10% Density, 10% hit rate + // Keys: 0, 10, 20, ... (very sparse), only 10% of probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_semi_d10_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 100% Density, ~1% hit rate, fanout ~100 + // Build keys are duplicated: 100K rows over 1K distinct keys. Matching + // probe rows produce many duplicate probe indices before RightSemi + // deduplication. + { + let fanout_keys = 1_000; + let left_batches = build_batches(build_rows, fanout_keys, 0, &s); + let right_batches = build_batches(probe_rows, build_rows, 0, &s); + group.bench_function( + BenchmarkId::new("right_semi_fanout100_h1", probe_rows), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }, + ); + } + + // ========================================================================= + // RightAnti Join benchmarks + // ========================================================================= + + // RightAnti - 100% Density, 100% hit rate (no output) + // Keys: 0..100K contiguous, all probe rows find a match -> no output + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows, 0, &s); + group.bench_function(BenchmarkId::new("right_anti_d100_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 100% Density, 10% hit rate (90% output) + // Keys: 0..100K contiguous, only 10% of probe rows find a match -> 90% output + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows * 10, 0, &s); + group.bench_function(BenchmarkId::new("right_anti_d100_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 50% Density, 100% hit rate (no output) + // Keys: 0, 2, 4, ... (sparse), all probe rows find a match -> no output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_anti_d50_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 50% Density, 10% hit rate (90% output) + // Keys: 0, 2, 4, ... (sparse), only 10% of probe rows find a match -> 90% output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_anti_d50_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 10% Density, 100% hit rate (no output) + // Keys: 0, 10, 20, ... (very sparse), all probe rows find a match -> no output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_anti_d10_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 10% Density, 10% hit rate (90% output) + // Keys: 0, 10, 20, ... (very sparse), only 10% of probe rows find a match -> 90% output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_anti_d10_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_hash_join_semi_anti); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/multi_group_by.rs b/native/vendor/datafusion-physical-plan/benches/multi_group_by.rs new file mode 100644 index 0000000..0c689f9 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/multi_group_by.rs
@@ -0,0 +1,815 @@ +// 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 multi-column GROUP BY performance comparing vectorized +//! (`GroupValuesColumn`) vs row-based (`GroupValuesRows`) implementations. +//! +//! Motivated by <https://github.com/apache/datafusion/issues/17850> which +//! showed vectorized can regress for low-cardinality, high-row-count scenarios. +//! +//! Uses the direct `GroupValues::intern()` API with identical data for both +//! implementations — a fair apples-to-apples comparison with the same hashing +//! and data layout. Most experiments use `Int32` columns; `bench_fixed_size_binary` +//! covers a `(FixedSizeBinary, Int32)` key to exercise the +//! `FixedSizeBinaryGroupValueBuilder`. + +use arrow::array::{ + ArrayRef, Decimal256Array, DurationMicrosecondArray, Float16Array, Int32Array, + IntervalMonthDayNanoArray, UInt32Array, +}; +use arrow::compute::take; +use arrow::datatypes::{ + DataType, Field, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, TimeUnit, + i256, +}; +use arrow::util::bench_util::create_fsb_array; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_physical_plan::aggregates::group_values::GroupValues; +use datafusion_physical_plan::aggregates::group_values::GroupValuesRows; +use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn; +use half::f16; +use std::hint::black_box; +use std::sync::Arc; + +const DEFAULT_BATCH_SIZE: usize = 8192; + +fn make_schema(num_cols: usize) -> SchemaRef { + let fields: Vec<Field> = (0..num_cols) + .map(|i| Field::new(format!("col_{i}"), DataType::Int32, false)) + .collect(); + Arc::new(Schema::new(fields)) +} + +fn generate_batches( + num_cols: usize, + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec<Vec<ArrayRef>> { + let per_col_card = (num_distinct_groups as f64) + .powf(1.0 / num_cols as f64) + .ceil() as usize; + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + (0..num_cols) + .map(|col_idx| { + let values: Vec<i32> = (0..current_batch_size) + .map(|row| { + let global_row = batch_start + row; + let group_id = global_row % num_distinct_groups; + let divisor = per_col_card.pow(col_idx as u32); + ((group_id / divisor) % per_col_card) as i32 + }) + .collect(); + Arc::new(Int32Array::from(values)) as ArrayRef + }) + .collect() + }) + .collect() +} + +fn create_group_values(schema: &SchemaRef, vectorized: bool) -> Box<dyn GroupValues> { + if vectorized { + Box::new(GroupValuesColumn::<false>::try_new(Arc::clone(schema)).unwrap()) + } else { + Box::new(GroupValuesRows::try_new(Arc::clone(schema)).unwrap()) + } +} + +fn bench_intern( + gv: &mut Box<dyn GroupValues>, + batches: &[Vec<ArrayRef>], + groups: &mut Vec<usize>, +) { + for batch in batches { + groups.clear(); + gv.intern(batch, groups).unwrap(); + } + black_box(&*groups); +} + +/// Experiment 1: Issue #17850 regression scenario. +/// 3 columns, 64 groups (4^3), scaling row count. +fn bench_issue_17850_regression(c: &mut Criterion) { + let mut group = c.benchmark_group("issue_17850_regression"); + group.sample_size(10); + + let num_cols = 3; + let num_groups = 64; + let schema = make_schema(num_cols); + + for num_rows in [1_000_000, 5_000_000, 10_000_000, 20_000_000, 50_000_000] { + let batches = + generate_batches(num_cols, num_groups, num_rows, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("{num_rows}_rows")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 2: Low cardinality sweep. +fn bench_low_cardinality(c: &mut Criterion) { + let mut group = c.benchmark_group("low_cardinality"); + group.sample_size(15); + + for (num_cols, per_col_card) in + [(3usize, 2usize), (3, 4), (3, 8), (4, 2), (4, 4), (4, 8)] + { + let num_groups = per_col_card.pow(num_cols as u32); + let schema = make_schema(num_cols); + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new( + label, + format!("cols_{num_cols}_card_{per_col_card}_grp_{num_groups}"), + ), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 3: Batch size sensitivity. +fn bench_batch_size_sensitivity(c: &mut Criterion) { + let mut group = c.benchmark_group("batch_size_sensitivity"); + group.sample_size(10); + + let num_cols = 3; + let num_groups = 64; + let schema = make_schema(num_cols); + + for batch_size in [1024, 4096, 8192, 16384, 32768] { + let batches = generate_batches(num_cols, num_groups, 1_000_000, batch_size); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("batch_{batch_size}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(batch_size), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 4: Column count scaling with low groups. +fn bench_column_scaling(c: &mut Criterion) { + let mut group = c.benchmark_group("column_scaling"); + group.sample_size(15); + + let cases: &[(usize, usize)] = + &[(2, 100), (3, 125), (4, 81), (6, 729), (8, 256), (10, 1024)]; + + for &(num_cols, num_groups) in cases { + let schema = make_schema(num_cols); + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("cols_{num_cols}_grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 5: High cardinality column scaling (~1M groups). +fn bench_high_cardinality_scaling(c: &mut Criterion) { + let mut group = c.benchmark_group("high_cardinality_scaling"); + group.sample_size(10); + + for num_cols in [2, 3, 4, 6, 8, 10] { + let num_groups = 1_000_000; + let schema = make_schema(num_cols); + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("cols_{num_cols}_grp_1M")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 6: Group count sweep with fixed 4 columns. +fn bench_group_count_sweep(c: &mut Criterion) { + let mut group = c.benchmark_group("group_count_sweep"); + group.sample_size(15); + + let num_cols = 4; + let schema = make_schema(num_cols); + + for num_groups in [ + 16, 64, 256, 1000, 5000, 10_000, 50_000, 100_000, 500_000, 1_000_000, + ] { + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Width in bytes of the FixedSizeBinary group column (UUID-sized). +const FSB_WIDTH: usize = 16; + +/// Schema for the FixedSizeBinary experiment: a `FixedSizeBinary` group column +/// paired with an `Int32` column, exercising a multi-column GROUP BY that +/// includes a fixed-width binary key (e.g. grouping on a UUID). +fn make_fsb_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("fsb", DataType::FixedSizeBinary(FSB_WIDTH as i32), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(FixedSizeBinary, Int32)` batches with exactly +/// `num_distinct_groups` distinct keys. +/// +/// The distinct FixedSizeBinary values come from arrow-rs's `create_fsb_array` +/// benchmark generator; rows cycle through that pool (mirroring how +/// `generate_batches` controls Int32 cardinality) so the group count is +/// controlled. The `Int32` column is keyed identically, keeping the combined +/// cardinality equal to `num_distinct_groups`. +fn generate_fsb_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec<Vec<ArrayRef>> { + // Pool of distinct FixedSizeBinary values (fixed seed, no nulls). + let pool = create_fsb_array(num_distinct_groups, 0.0, FSB_WIDTH); + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let indices: UInt32Array = group_ids.clone().map(|g| g as u32).collect(); + let fsb = take(&pool, &indices, None).unwrap(); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![fsb, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 7: Group count sweep for a `(FixedSizeBinary, Int32)` key. +/// +/// Exercises the `FixedSizeBinaryGroupValueBuilder` used by multi-column +/// GROUP BY. Before FixedSizeBinary support, such a schema fell back to the +/// row-based `GroupValuesRows`; this compares the vectorized columnar path +/// (`vectorized`) against that baseline (`row_based`). +fn bench_fixed_size_binary(c: &mut Criterion) { + let mut group = c.benchmark_group("fixed_size_binary"); + group.sample_size(15); + + let schema = make_fsb_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = generate_fsb_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_f16_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("f16", DataType::Float16, false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Float16, Int32)` batches with `num_distinct_groups` distinct keys. +/// +/// `f16` has only ~63.5k finite values, so `num_distinct_groups` must stay well +/// under that (see `bench_float16`). Distinct keys are the low finite `f16` bit +/// patterns, skipping NaN and inf. The `Int32` column is keyed identically so +/// the combined cardinality equals `num_distinct_groups`. +fn generate_f16_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec<Vec<ArrayRef>> { + let pool: Vec<f16> = (0u16..) + .map(f16::from_bits) + .filter(|v| v.is_finite()) + .take(num_distinct_groups) + .collect(); + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = Float16Array::from_iter_values(group_ids.clone().map(|g| pool[g])); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 8: Group count sweep for a `(Float16, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Float16` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +/// Group counts are capped below `f16`'s ~63.5k distinct finite values. +fn bench_float16(c: &mut Criterion) { + let mut group = c.benchmark_group("float16"); + group.sample_size(15); + + let schema = make_f16_schema(); + + for num_groups in [1_000, 60_000] { + let batches = generate_f16_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_duration_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("dur", DataType::Duration(TimeUnit::Microsecond), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Duration(Microsecond), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct duration is `g` microseconds. The `Int32` column is keyed +/// identically so the combined cardinality equals `num_distinct_groups`. +fn generate_duration_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec<Vec<ArrayRef>> { + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = DurationMicrosecondArray::from_iter_values( + group_ids.clone().map(|g| g as i64), + ); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 9: Group count sweep for a `(Duration, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Duration` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +fn bench_duration(c: &mut Criterion) { + let mut group = c.benchmark_group("duration"); + group.sample_size(15); + + let schema = make_duration_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_duration_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_interval_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("iv", DataType::Interval(IntervalUnit::MonthDayNano), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Interval(MonthDayNano), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct interval is `MonthDayNano(g, 0, 0)`. The `Int32` column is +/// keyed identically so the combined cardinality equals `num_distinct_groups`. +fn generate_interval_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec<Vec<ArrayRef>> { + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = IntervalMonthDayNanoArray::from_iter_values( + group_ids + .clone() + .map(|g| IntervalMonthDayNano::new(g as i32, 0, 0)), + ); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 10: Group count sweep for an `(Interval, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Interval` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +fn bench_interval(c: &mut Criterion) { + let mut group = c.benchmark_group("interval"); + group.sample_size(15); + + let schema = make_interval_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_interval_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_decimal256_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("dec", DataType::Decimal256(50, 0), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Decimal256(50, 0), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct value is `i256::from_i128(g)`, and precision > 38 keeps it a +/// genuine `Decimal256`. The `Int32` column is keyed identically so the combined +/// cardinality equals `num_distinct_groups`. +fn generate_decimal256_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec<Vec<ArrayRef>> { + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = Decimal256Array::from_iter_values( + group_ids.clone().map(|g| i256::from_i128(g as i128)), + ) + .with_precision_and_scale(50, 0) + .unwrap(); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 11: Group count sweep for a `(Decimal256, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Decimal256` (32-byte +/// `i256` native) on the multi-column path (previously such a schema fell back +/// to `GroupValuesRows`). +fn bench_decimal256(c: &mut Criterion) { + let mut group = c.benchmark_group("decimal256"); + group.sample_size(15); + + let schema = make_decimal256_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_decimal256_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::<usize>::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_issue_17850_regression, + bench_low_cardinality, + bench_batch_size_sensitivity, + bench_column_scaling, + bench_high_cardinality_scaling, + bench_group_count_sweep, + bench_fixed_size_binary, + bench_float16, + bench_duration, + bench_interval, + bench_decimal256, +); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/partial_ordering.rs b/native/vendor/datafusion-physical-plan/benches/partial_ordering.rs new file mode 100644 index 0000000..bdadd62 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/partial_ordering.rs
@@ -0,0 +1,60 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array}; +use datafusion_physical_plan::aggregates::order::GroupOrderingPartial; + +use criterion::{Criterion, criterion_group, criterion_main}; + +const BATCH_SIZE: usize = 8192; + +fn create_test_arrays(num_columns: usize) -> Vec<ArrayRef> { + (0..num_columns) + .map(|i| { + Arc::new(Int32Array::from_iter_values( + (0..BATCH_SIZE as i32).map(|x| x * (i + 1) as i32), + )) as ArrayRef + }) + .collect() +} +fn bench_new_groups(c: &mut Criterion) { + let mut group = c.benchmark_group("group_ordering_partial"); + + // Test with 1, 2, 4, and 8 order indices + for num_columns in [1, 2, 4, 8] { + let order_indices: Vec<usize> = (0..num_columns).collect(); + + group.bench_function(format!("order_indices_{num_columns}"), |b| { + let batch_group_values = create_test_arrays(num_columns); + let group_indices: Vec<usize> = (0..BATCH_SIZE).collect(); + + b.iter(|| { + let mut ordering = + GroupOrderingPartial::try_new(order_indices.clone()).unwrap(); + ordering + .new_groups(&batch_group_values, &group_indices, BATCH_SIZE) + .unwrap(); + }); + }); + } + group.finish(); +} + +criterion_group!(benches, bench_new_groups); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/sort_merge_join.rs b/native/vendor/datafusion-physical-plan/benches/sort_merge_join.rs new file mode 100644 index 0000000..82610b2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/sort_merge_join.rs
@@ -0,0 +1,204 @@ +// 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. + +//! Criterion benchmarks for Sort Merge Join +//! +//! These benchmarks measure the join kernel in isolation by feeding +//! pre-sorted RecordBatches directly into SortMergeJoinExec, avoiding +//! sort / scan overhead. + +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch, StringArray}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::NullEquality; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::collect; +use datafusion_physical_plan::joins::{SortMergeJoinExec, utils::JoinOn}; +use datafusion_physical_plan::test::TestMemoryExec; +use tokio::runtime::Runtime; + +/// Build pre-sorted RecordBatches (split into ~8192-row chunks). +/// +/// Schema: (key: Int64, data: Int64, payload: Utf8) +/// +/// `key_mod` controls distinct key count: key = row_index % key_mod. +fn build_sorted_batches( + num_rows: usize, + key_mod: usize, + schema: &SchemaRef, +) -> Vec<RecordBatch> { + let mut rows: Vec<(i64, i64)> = (0..num_rows) + .map(|i| ((i % key_mod) as i64, i as i64)) + .collect(); + rows.sort(); + + let keys: Vec<i64> = rows.iter().map(|(k, _)| *k).collect(); + let data: Vec<i64> = rows.iter().map(|(_, d)| *d).collect(); + let payload: Vec<String> = data.iter().map(|d| format!("val_{d}")).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from(data)), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn make_exec( + batches: &[RecordBatch], + schema: &SchemaRef, +) -> Arc<dyn datafusion_physical_plan::ExecutionPlan> { + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("data", DataType::Int64, false), + Field::new("payload", DataType::Utf8, false), + ])) +} + +fn do_join( + left: Arc<dyn datafusion_physical_plan::ExecutionPlan>, + right: Arc<dyn datafusion_physical_plan::ExecutionPlan>, + join_type: datafusion_common::JoinType, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let join = SortMergeJoinExec::try_new( + left, + right, + on, + None, + join_type, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let task_ctx = Arc::new(TaskContext::default()); + rt.block_on(async { + let batches = collect(Arc::new(join), task_ctx).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }) +} + +fn bench_smj(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + + let mut group = c.benchmark_group("sort_merge_join"); + + // 1:1 Inner Join — 100K rows each, unique keys + // Best case for contiguous-range optimization: every index array is [0,1,2,...]. + { + let n = 100_000; + let left_batches = build_sorted_batches(n, n, &s); + let right_batches = build_sorted_batches(n, n, &s); + group.bench_function(BenchmarkId::new("inner_1to1", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::Inner, &rt) + }) + }); + } + + // 1:10 Inner Join — 100K left, 100K right, 10K distinct keys + { + let n = 100_000; + let key_mod = 10_000; + let left_batches = build_sorted_batches(n, key_mod, &s); + let right_batches = build_sorted_batches(n, key_mod, &s); + group.bench_function(BenchmarkId::new("inner_1to10", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::Inner, &rt) + }) + }); + } + + // Left Join — 100K each, ~5% unmatched on left + { + let n = 100_000; + let left_batches = build_sorted_batches(n, n + n / 20, &s); + let right_batches = build_sorted_batches(n, n, &s); + group.bench_function(BenchmarkId::new("left_1to1_unmatched", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::Left, &rt) + }) + }); + } + + // Left Semi Join — 100K left, 100K right, 10K keys + { + let n = 100_000; + let key_mod = 10_000; + let left_batches = build_sorted_batches(n, key_mod, &s); + let right_batches = build_sorted_batches(n, key_mod, &s); + group.bench_function(BenchmarkId::new("left_semi_1to10", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::LeftSemi, &rt) + }) + }); + } + + // Left Anti Join — 100K left, 100K right, partial match + { + let n = 100_000; + let left_batches = build_sorted_batches(n, n + n / 5, &s); + let right_batches = build_sorted_batches(n, n, &s); + group.bench_function(BenchmarkId::new("left_anti_partial", n), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_join(left, right, datafusion_common::JoinType::LeftAnti, &rt) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_smj); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/sort_preserving_merge.rs b/native/vendor/datafusion-physical-plan/benches/sort_preserving_merge.rs new file mode 100644 index 0000000..76ebf23 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/sort_preserving_merge.rs
@@ -0,0 +1,197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::{ + array::{ArrayRef, StringArray, UInt64Array}, + record_batch::RecordBatch, +}; +use arrow_schema::{SchemaRef, SortOptions}; +use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ + collect, sorts::sort_preserving_merge::SortPreservingMergeExec, +}; + +use std::sync::Arc; + +const BENCH_ROWS: usize = 1_000_000; // 1 million rows + +fn get_large_string(idx: usize) -> String { + let base_content = [ + concat!( + "# Advanced Topics in Computer Science\n\n", + "## Summary\nThis article explores complex system design patterns and...\n\n", + "```rust\nfn process_data(data: &mut [i32]) {\n // Parallel processing example\n data.par_iter_mut().for_each(|x| *x *= 2);\n}\n```\n\n", + "## Performance Considerations\nWhen implementing concurrent systems...\n" + ), + concat!( + "## API Documentation\n\n", + "```json\n{\n \"endpoint\": \"/api/v2/users\",\n \"methods\": [\"GET\", \"POST\"],\n \"parameters\": {\n \"page\": \"number\"\n }\n}\n```\n\n", + "# Authentication Guide\nSecure your API access using OAuth 2.0...\n" + ), + concat!( + "# Data Processing Pipeline\n\n", + "```python\nfrom multiprocessing import Pool\n\ndef main():\n with Pool(8) as p:\n results = p.map(process_item, data)\n```\n\n", + "## Summary of Optimizations\n1. Batch processing\n2. Memory pooling\n3. Concurrent I/O operations\n" + ), + concat!( + "# System Architecture Overview\n\n", + "## Components\n- Load Balancer\n- Database Cluster\n- Cache Service\n\n", + "```go\nfunc main() {\n router := gin.Default()\n router.GET(\"/api/health\", healthCheck)\n router.Run(\":8080\")\n}\n```\n" + ), + concat!( + "## Configuration Reference\n\n", + "```yaml\nserver:\n port: 8080\n max_threads: 32\n\ndatabase:\n url: postgres://user@prod-db:5432/main\n```\n\n", + "# Deployment Strategies\nBlue-green deployment patterns with...\n" + ), + ]; + base_content[idx % base_content.len()].to_string() +} + +fn generate_sorted_string_column(rows: usize) -> ArrayRef { + let mut values = Vec::with_capacity(rows); + for i in 0..rows { + values.push(get_large_string(i)); + } + values.sort(); + Arc::new(StringArray::from(values)) +} + +fn generate_sorted_u64_column(rows: usize) -> ArrayRef { + Arc::new(UInt64Array::from((0_u64..rows as u64).collect::<Vec<_>>())) +} + +fn create_partitions<const IS_LARGE_COLUMN_TYPE: bool>( + num_partitions: usize, + num_columns: usize, + num_rows: usize, +) -> Vec<Vec<RecordBatch>> { + (0..num_partitions) + .map(|_| { + let rows = (0..num_columns) + .map(|i| { + ( + format!("col-{i}"), + if IS_LARGE_COLUMN_TYPE { + generate_sorted_string_column(num_rows) + } else { + generate_sorted_u64_column(num_rows) + }, + ) + }) + .collect::<Vec<_>>(); + + let batch = RecordBatch::try_from_iter(rows).unwrap(); + vec![batch] + }) + .collect() +} + +struct BenchData { + bench_name: String, + partitions: Vec<Vec<RecordBatch>>, + schema: SchemaRef, + sort_order: LexOrdering, +} + +fn get_bench_data() -> Vec<BenchData> { + let mut ret = Vec::new(); + let mut push_bench_data = |bench_name: &str, partitions: Vec<Vec<RecordBatch>>| { + let schema = partitions[0][0].schema(); + // Define sort order (col1 ASC, col2 ASC, col3 ASC) + let sort_order = LexOrdering::new(schema.fields().iter().map(|field| { + PhysicalSortExpr::new( + col(field.name(), &schema).unwrap(), + SortOptions::default(), + ) + })) + .unwrap(); + ret.push(BenchData { + bench_name: bench_name.to_string(), + partitions, + schema, + sort_order, + }); + }; + // 1. single large string column + { + let partitions = create_partitions::<true>(3, 1, BENCH_ROWS); + push_bench_data("single_large_string_column_with_1m_rows", partitions); + } + // 2. single u64 column + { + let partitions = create_partitions::<false>(3, 1, BENCH_ROWS); + push_bench_data("single_u64_column_with_1m_rows", partitions); + } + // 3. multiple large string columns + { + let partitions = create_partitions::<true>(3, 3, BENCH_ROWS); + push_bench_data("multiple_large_string_columns_with_1m_rows", partitions); + } + // 4. multiple u64 columns + { + let partitions = create_partitions::<false>(3, 3, BENCH_ROWS); + push_bench_data("multiple_u64_columns_with_1m_rows", partitions); + } + ret +} + +/// Add a benchmark to test the optimization effect of reusing Rows. +/// Run this benchmark with: +/// ```sh +/// cargo bench --features="bench" --bench sort_preserving_merge -- --sample-size=10 +/// ``` +fn bench_merge_sorted_preserving(c: &mut Criterion) { + let task_ctx = Arc::new(TaskContext::default()); + let bench_data = get_bench_data(); + for data in bench_data.into_iter() { + let BenchData { + bench_name, + partitions, + schema, + sort_order, + } = data; + c.bench_function( + &format!("bench_merge_sorted_preserving/{bench_name}"), + |b| { + b.iter_batched( + || { + let exec = TestMemoryExec::try_new_exec( + &partitions, + schema.clone(), + None, + ) + .unwrap(); + Arc::new(SortPreservingMergeExec::new(sort_order.clone(), exec)) + }, + |merge_exec| { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + collect(merge_exec, task_ctx.clone()).await.unwrap(); + }); + }, + BatchSize::LargeInput, + ) + }, + ); + } +} + +criterion_group!(benches, bench_merge_sorted_preserving); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/benches/spill_io.rs b/native/vendor/datafusion-physical-plan/benches/spill_io.rs new file mode 100644 index 0000000..ddd83ca --- /dev/null +++ b/native/vendor/datafusion-physical-plan/benches/spill_io.rs
@@ -0,0 +1,581 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ + Date32Builder, Decimal128Builder, Int32Builder, Int64Builder, RecordBatch, + StringBuilder, +}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::measurement::WallTime; +use criterion::{ + BatchSize, BenchmarkGroup, BenchmarkId, Criterion, criterion_group, criterion_main, +}; +use datafusion_common::config::SpillCompression; +use datafusion_common::human_readable_size; +use datafusion_common::instant::Instant; +use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_physical_plan::SpillManager; +use datafusion_physical_plan::common::collect; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; +use rand::{Rng, SeedableRng}; +use std::sync::Arc; +use tokio::runtime::Runtime; + +pub fn create_batch(num_rows: usize, allow_nulls: bool) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Utf8, true), + Field::new("c2", DataType::Date32, true), + Field::new("c3", DataType::Decimal128(11, 2), true), + ])); + + let mut a = Int32Builder::new(); + let mut b = StringBuilder::new(); + let mut c = Date32Builder::new(); + let mut d = Decimal128Builder::new() + .with_precision_and_scale(11, 2) + .unwrap(); + + for i in 0..num_rows { + a.append_value(i as i32); + c.append_value(i as i32); + d.append_value((i * 1000000) as i128); + if allow_nulls && i % 10 == 0 { + b.append_null(); + } else { + b.append_value(format!("this is string number {i}")); + } + } + + let a = a.finish(); + let b = b.finish(); + let c = c.finish(); + let d = d.finish(); + + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(a), Arc::new(b), Arc::new(c), Arc::new(d)], + ) + .unwrap() +} + +// BENCHMARK: REVALIDATION OVERHEAD COMPARISON +// --------------------------------------------------------- +// To compare performance with/without Arrow IPC validation: +// +// 1. Locate the function `read_spill` +// 2. Modify the `skip_validation` flag: +// - Set to `false` to enable validation +// 3. Rerun `cargo bench --bench spill_io` +fn bench_spill_io(c: &mut Criterion) { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = Arc::new(Schema::new(vec![ + Field::new("c0", DataType::Int32, true), + Field::new("c1", DataType::Utf8, true), + Field::new("c2", DataType::Date32, true), + Field::new("c3", DataType::Decimal128(11, 2), true), + ])); + let spill_manager = SpillManager::new(env, metrics, schema); + + let mut group = c.benchmark_group("spill_io"); + let rt = Runtime::new().unwrap(); + + group.bench_with_input( + BenchmarkId::new("StreamReader/read_100", ""), + &spill_manager, + |b, spill_manager| { + b.iter_batched( + // Setup phase: Create fresh state for each benchmark iteration. + // - generate an ipc file. + // This ensures each iteration starts with clean resources. + || { + let batch = create_batch(8192, true); + spill_manager + .spill_record_batch_and_finish(&vec![batch; 100], "Test") + .unwrap() + .unwrap() + }, + // Benchmark phase: + // - Execute the read operation via SpillManager + // - Wait for the consumer to finish processing + |spill_file| { + rt.block_on(async { + let stream = spill_manager + .read_spill_as_stream(spill_file, None) + .unwrap(); + let _ = collect(stream).await.unwrap(); + }) + }, + BatchSize::LargeInput, + ) + }, + ); + group.finish(); +} + +// Generate `num_batches` RecordBatches mimicking TPC-H Q2's partial aggregate result: +// GROUP BY ps_partkey -> MIN(ps_supplycost) +fn create_q2_like_batches( + num_batches: usize, + num_rows: usize, +) -> (Arc<Schema>, Vec<RecordBatch>) { + // use fixed seed + let seed = 2; + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut batches = Vec::with_capacity(num_batches); + + let mut current_key = 400000_i64; + + let schema = Arc::new(Schema::new(vec![ + Field::new("ps_partkey", DataType::Int64, false), + Field::new("min_ps_supplycost", DataType::Decimal128(15, 2), true), + ])); + + for _ in 0..num_batches { + let mut partkey_builder = Int64Builder::new(); + let mut cost_builder = Decimal128Builder::new() + .with_precision_and_scale(15, 2) + .unwrap(); + + for _ in 0..num_rows { + // Occasionally skip a few partkey values to simulate sparsity + let jump = if rng.random_bool(0.05) { + rng.random_range(2..10) + } else { + 1 + }; + current_key += jump; + + let supply_cost = rng.random_range(10_00..100_000) as i128; + + partkey_builder.append_value(current_key); + cost_builder.append_value(supply_cost); + } + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(partkey_builder.finish()), + Arc::new(cost_builder.finish()), + ], + ) + .unwrap(); + + batches.push(batch); + } + + (schema, batches) +} + +/// Generate `num_batches` RecordBatches mimicking TPC-H Q16's partial aggregate result: +/// GROUP BY (p_brand, p_type, p_size) -> COUNT(DISTINCT ps_suppkey) +pub fn create_q16_like_batches( + num_batches: usize, + num_rows: usize, +) -> (Arc<Schema>, Vec<RecordBatch>) { + let seed = 16; + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut batches = Vec::with_capacity(num_batches); + + let schema = Arc::new(Schema::new(vec![ + Field::new("p_brand", DataType::Utf8, false), + Field::new("p_type", DataType::Utf8, false), + Field::new("p_size", DataType::Int32, false), + Field::new("alias1", DataType::Int64, false), // COUNT(DISTINCT ps_suppkey) + ])); + + // Representative string pools + let brands = ["Brand#32", "Brand#33", "Brand#41", "Brand#42", "Brand#55"]; + let types = [ + "PROMO ANODIZED NICKEL", + "STANDARD BRUSHED NICKEL", + "PROMO POLISHED COPPER", + "ECONOMY ANODIZED BRASS", + "LARGE BURNISHED COPPER", + "STANDARD POLISHED TIN", + "SMALL PLATED STEEL", + "MEDIUM POLISHED COPPER", + ]; + let sizes = [3, 9, 14, 19, 23, 36, 45, 49]; + + for _ in 0..num_batches { + let mut brand_builder = StringBuilder::new(); + let mut type_builder = StringBuilder::new(); + let mut size_builder = Int32Builder::new(); + let mut count_builder = Int64Builder::new(); + + for _ in 0..num_rows { + let brand = brands[rng.random_range(0..brands.len())]; + let ptype = types[rng.random_range(0..types.len())]; + let size = sizes[rng.random_range(0..sizes.len())]; + let count = rng.random_range(1000..100_000); + + brand_builder.append_value(brand); + type_builder.append_value(ptype); + size_builder.append_value(size); + count_builder.append_value(count); + } + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(brand_builder.finish()), + Arc::new(type_builder.finish()), + Arc::new(size_builder.finish()), + Arc::new(count_builder.finish()), + ], + ) + .unwrap(); + + batches.push(batch); + } + + (schema, batches) +} + +// Generate `num_batches` RecordBatches mimicking TPC-H Q20's partial aggregate result: +// GROUP BY (l_partkey, l_suppkey) -> SUM(l_quantity) +fn create_q20_like_batches( + num_batches: usize, + num_rows: usize, +) -> (Arc<Schema>, Vec<RecordBatch>) { + let seed = 20; + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut batches = Vec::with_capacity(num_batches); + + let mut current_partkey = 400000_i64; + + let schema = Arc::new(Schema::new(vec![ + Field::new("l_partkey", DataType::Int64, false), + Field::new("l_suppkey", DataType::Int64, false), + Field::new("sum_l_quantity", DataType::Decimal128(25, 2), true), + ])); + + for _ in 0..num_batches { + let mut partkey_builder = Int64Builder::new(); + let mut suppkey_builder = Int64Builder::new(); + let mut quantity_builder = Decimal128Builder::new() + .with_precision_and_scale(25, 2) + .unwrap(); + + for _ in 0..num_rows { + // Occasionally skip a few partkey values to simulate sparsity + let partkey_jump = if rng.random_bool(0.03) { + rng.random_range(2..6) + } else { + 1 + }; + current_partkey += partkey_jump; + + let suppkey = rng.random_range(10_000..99_999); + let quantity = rng.random_range(500..20_000) as i128; + + partkey_builder.append_value(current_partkey); + suppkey_builder.append_value(suppkey); + quantity_builder.append_value(quantity); + } + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(partkey_builder.finish()), + Arc::new(suppkey_builder.finish()), + Arc::new(quantity_builder.finish()), + ], + ) + .unwrap(); + + batches.push(batch); + } + + (schema, batches) +} + +/// Generate `num_batches` wide RecordBatches resembling sort-tpch Q10 for benchmarking. +/// This includes multiple numeric, date, and Utf8View columns (15 total). +pub fn create_wide_batches( + num_batches: usize, + num_rows: usize, +) -> (Arc<Schema>, Vec<RecordBatch>) { + let seed = 10; + let mut rng = rand::rngs::StdRng::seed_from_u64(seed); + let mut batches = Vec::with_capacity(num_batches); + + let schema = Arc::new(Schema::new(vec![ + Field::new("l_linenumber", DataType::Int32, false), + Field::new("l_suppkey", DataType::Int64, false), + Field::new("l_orderkey", DataType::Int64, false), + Field::new("l_partkey", DataType::Int64, false), + Field::new("l_quantity", DataType::Decimal128(15, 2), false), + Field::new("l_extendedprice", DataType::Decimal128(15, 2), false), + Field::new("l_discount", DataType::Decimal128(15, 2), false), + Field::new("l_tax", DataType::Decimal128(15, 2), false), + Field::new("l_returnflag", DataType::Utf8, false), + Field::new("l_linestatus", DataType::Utf8, false), + Field::new("l_shipdate", DataType::Date32, false), + Field::new("l_commitdate", DataType::Date32, false), + Field::new("l_receiptdate", DataType::Date32, false), + Field::new("l_shipinstruct", DataType::Utf8, false), + Field::new("l_shipmode", DataType::Utf8, false), + ])); + + for _ in 0..num_batches { + let mut linenum = Int32Builder::new(); + let mut suppkey = Int64Builder::new(); + let mut orderkey = Int64Builder::new(); + let mut partkey = Int64Builder::new(); + let mut quantity = Decimal128Builder::new() + .with_precision_and_scale(15, 2) + .unwrap(); + let mut extprice = Decimal128Builder::new() + .with_precision_and_scale(15, 2) + .unwrap(); + let mut discount = Decimal128Builder::new() + .with_precision_and_scale(15, 2) + .unwrap(); + let mut tax = Decimal128Builder::new() + .with_precision_and_scale(15, 2) + .unwrap(); + let mut retflag = StringBuilder::new(); + let mut linestatus = StringBuilder::new(); + let mut shipdate = Date32Builder::new(); + let mut commitdate = Date32Builder::new(); + let mut receiptdate = Date32Builder::new(); + let mut shipinstruct = StringBuilder::new(); + let mut shipmode = StringBuilder::new(); + + let return_flags = ["A", "N", "R"]; + let statuses = ["F", "O"]; + let instructs = ["DELIVER IN PERSON", "COLLECT COD", "NONE"]; + let modes = ["TRUCK", "MAIL", "SHIP", "RAIL", "AIR"]; + + for i in 0..num_rows { + linenum.append_value((i % 7) as i32); + suppkey.append_value(rng.random_range(0..100_000)); + orderkey.append_value(1_000_000 + i as i64); + partkey.append_value(rng.random_range(0..200_000)); + + quantity.append_value(rng.random_range(100..10000) as i128); + extprice.append_value(rng.random_range(1_000..1_000_000) as i128); + discount.append_value(rng.random_range(0..10000) as i128); + tax.append_value(rng.random_range(0..5000) as i128); + + retflag.append_value(return_flags[rng.random_range(0..return_flags.len())]); + linestatus.append_value(statuses[rng.random_range(0..statuses.len())]); + + let base_date = 10_000; + shipdate.append_value(base_date + (i % 1000) as i32); + commitdate.append_value(base_date + (i % 1000) as i32 + 1); + receiptdate.append_value(base_date + (i % 1000) as i32 + 2); + + shipinstruct.append_value(instructs[rng.random_range(0..instructs.len())]); + shipmode.append_value(modes[rng.random_range(0..modes.len())]); + } + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(linenum.finish()), + Arc::new(suppkey.finish()), + Arc::new(orderkey.finish()), + Arc::new(partkey.finish()), + Arc::new(quantity.finish()), + Arc::new(extprice.finish()), + Arc::new(discount.finish()), + Arc::new(tax.finish()), + Arc::new(retflag.finish()), + Arc::new(linestatus.finish()), + Arc::new(shipdate.finish()), + Arc::new(commitdate.finish()), + Arc::new(receiptdate.finish()), + Arc::new(shipinstruct.finish()), + Arc::new(shipmode.finish()), + ], + ) + .unwrap(); + batches.push(batch); + } + (schema, batches) +} + +// Benchmarks spill write + read performance across multiple compression codecs +// using realistic input data inspired by TPC-H aggregate spill scenarios. +// +// This function prepares synthetic RecordBatches that mimic the schema and distribution +// of intermediate aggregate results from representative TPC-H queries (Q2, Q16, Q20) and sort-tpch Q10. +// The schemas of these batches are: +// Q2 [Int64, Decimal128] +// Q16 [Utf8, Utf8, Int32, Int64] +// Q20 [Int64, Int64, Decimal128] +// sort-tpch Q10 (wide batch) [Int32, Int64 * 3, Decimal128 * 4, Date * 3, Utf8 * 4] +// For each dataset: +// - It evaluates spill performance under different compression codecs (e.g., Uncompressed, Zstd, LZ4). +// - It measures end-to-end spill write + read performance using Criterion. +// - It prints the observed memory-to-disk compression ratio for each codec. +// +// This helps evaluate the tradeoffs between compression ratio and runtime overhead for various codecs. +fn bench_spill_compression(c: &mut Criterion) { + let env = Arc::new(RuntimeEnv::default()); + let mut group = c.benchmark_group("spill_compression"); + let rt = Runtime::new().unwrap(); + let compressions = vec![ + SpillCompression::Uncompressed, + SpillCompression::Zstd, + SpillCompression::Lz4Frame, + ]; + + // Modify these values to change data volume. Note that each batch contains `num_rows` rows. + let num_batches = 50; + let num_rows = 8192; + + // Q2 [Int64, Decimal128] + let (schema, batches) = create_q2_like_batches(num_batches, num_rows); + benchmark_spill_batches_for_all_codec( + &mut group, + "q2", + batches, + &compressions, + &rt, + env.clone(), + schema, + ); + // Q16 [Utf8, Utf8, Int32, Int64] + let (schema, batches) = create_q16_like_batches(num_batches, num_rows); + benchmark_spill_batches_for_all_codec( + &mut group, + "q16", + batches, + &compressions, + &rt, + env.clone(), + schema, + ); + // Q20 [Int64, Int64, Decimal128] + let (schema, batches) = create_q20_like_batches(num_batches, num_rows); + benchmark_spill_batches_for_all_codec( + &mut group, + "q20", + batches, + &compressions, + &rt, + env.clone(), + schema, + ); + // sort-tpch Q10 (wide batch) [Int32, Int64 * 3, Decimal128 * 4, Date * 3, Utf8 * 4] + let (schema, batches) = create_wide_batches(num_batches, num_rows); + benchmark_spill_batches_for_all_codec( + &mut group, + "wide", + batches, + &compressions, + &rt, + env, + schema, + ); + group.finish(); +} + +#[expect(clippy::needless_pass_by_value)] +fn benchmark_spill_batches_for_all_codec( + group: &mut BenchmarkGroup<'_, WallTime>, + batch_label: &str, + batches: Vec<RecordBatch>, + compressions: &[SpillCompression], + rt: &Runtime, + env: Arc<RuntimeEnv>, + schema: Arc<Schema>, +) { + let mem_bytes: usize = batches.iter().map(|b| b.get_array_memory_size()).sum(); + + for &compression in compressions { + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = + SpillManager::new(Arc::clone(&env), metrics.clone(), Arc::clone(&schema)) + .with_compression_type(compression); + + let bench_id = BenchmarkId::new(batch_label, compression.to_string()); + group.bench_with_input(bench_id, &spill_manager, |b, spill_manager| { + b.iter_batched( + || batches.clone(), + |batches| { + rt.block_on(async { + let spill_file = spill_manager + .spill_record_batch_and_finish( + &batches, + &format!("{batch_label}_{compression}"), + ) + .unwrap() + .unwrap(); + let stream = spill_manager + .read_spill_as_stream(spill_file, None) + .unwrap(); + let _ = collect(stream).await.unwrap(); + }) + }, + BatchSize::LargeInput, + ) + }); + + // Run Spilling Read & Write once more to read file size & calculate bandwidth + let start = Instant::now(); + + let spill_file = spill_manager + .spill_record_batch_and_finish( + &batches, + &format!("{batch_label}_{compression}"), + ) + .unwrap() + .unwrap(); + + // calculate write_throughput (includes both compression and I/O time) based on in memory batch size + let write_time = start.elapsed(); + let write_throughput = (mem_bytes as u128 / write_time.as_millis().max(1)) * 1000; + + // calculate compression ratio + let disk_bytes = std::fs::metadata(spill_file.path().unwrap()) + .expect("metadata read fail") + .len() as usize; + let ratio = mem_bytes as f64 / disk_bytes.max(1) as f64; + + // calculate read_throughput (includes both compression and I/O time) based on in memory batch size + let rt = Runtime::new().unwrap(); + let start = Instant::now(); + rt.block_on(async { + let stream = spill_manager + .read_spill_as_stream(spill_file, None) + .unwrap(); + let _ = collect(stream).await.unwrap(); + }); + let read_time = start.elapsed(); + let read_throughput = (mem_bytes as u128 / read_time.as_millis().max(1)) * 1000; + + println!( + "[{} | {:?}] mem: {}| disk: {}| compression ratio: {:.3}x| throughput: (w) {}/s (r) {}/s", + batch_label, + compression, + human_readable_size(mem_bytes), + human_readable_size(disk_bytes), + ratio, + human_readable_size(write_throughput as usize), + human_readable_size(read_throughput as usize), + ); + } +} + +criterion_group!(benches, bench_spill_io, bench_spill_compression); +criterion_main!(benches);
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs new file mode 100644 index 0000000..91e9d65 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common.rs
@@ -0,0 +1,693 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_physical_expr::aggregate::AggregateFunctionExpr; + +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{ + AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, +}; + +/// Marker for raw rows -> partial state aggregation. +pub(in crate::aggregates) struct PartialMarker; +/// Marker for raw rows -> final value aggregation. +pub(in crate::aggregates) struct SingleMarker; +/// Marker for partial state -> partial state aggregation. +pub(in crate::aggregates) struct PartialReduceMarker; +/// Marker for raw rows -> partial state conversion without aggregation. +pub(in crate::aggregates) struct PartialSkipMarker; +/// Marker for partial state -> final value aggregation. +pub(in crate::aggregates) struct FinalMarker; + +/// Grouped hash table shared by the partial and final paths. +/// +/// While building, it consumes input batches and updates group / accumulator +/// state. While outputting, it incrementally drains that state into output +/// batches. +/// +/// # Logical and Physical Model +/// +/// Logically, this is a hash table that maps { group keys -> accumulator states } +/// For example, `AVG(v) GROUP BY k` stores one entry per `k`, where each +/// entry owns the `sum(v)` and `count(v)` state needed to compute the final +/// average. +/// +/// Physically, the group keys and accumulators are backed by [`GroupValues`] and +/// [`GroupsAccumulator`]. Both use columnar storage so aggregation can stay +/// vectorized. +/// +/// # Marker Type +/// `AggrMode` selects the aggregate semantics. +/// +/// e.g. `AggregateHashTable::<PartialMarker>::new(...)` creates an aggregate hash table +/// for the partial hash aggregate stage, the input schema is raw rows and output +/// schema is intermediate states. +/// +/// It is a zero-sized compile-time marker, so each stage keeps its update logic +/// in a separate impl block, to make the behavior difference explicit. +pub(in crate::aggregates) struct AggregateHashTable<AggrMode> { + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Raw input schema, used to evaluate expressions and synthesize empty + /// grouping-set rows. + pub(super) input_schema: SchemaRef, + + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Intermediate-state schema used when memory pressure requires the table + /// to spill its current state. + pub(super) state_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Lifecycle-specific state: building stage / outputting stage. + pub(super) state: AggregateHashTableState, + + pub(super) _mode: PhantomData<AggrMode>, +} + +/// Methods shared by all aggregate hash table modes. +impl<AggrMode> AggregateHashTable<AggrMode> { + pub(super) fn new_with_filters( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + filters: Vec<Option<Arc<dyn PhysicalExpr>>>, + ) -> Result<Self> { + if batch_size == 0 { + return internal_err!("AggregateHashTable requires config batch_size >= 1"); + } + + let input_schema = agg.input().schema(); + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + &agg.mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators: Vec<_> = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(HashAggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::<Result<_>>()?; + + let group_schema = agg.group_by.group_schema(&input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + + Ok(Self { + group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + input_schema, + output_schema, + state_schema, + batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + _mode: PhantomData, + }) + } + + /// See comments in [`EvaluatedAggregateBatch`] + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result<EvaluatedAggregateBatch> { + let state = self.state.building(); + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + // outer vec: one per each grouping set + // inner vec: all group by exprs for the current grouping set + let grouping_set_args = evaluate_group_by(&state.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + // The evaluated args for each accumulator + let accumulator_args = self + .state + .building() + .accumulators + .iter() + .map(|acc| acc.evaluate_acc_args(batch)) + .collect::<Result<Vec<_>>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Aggregates one input batch after selecting the mode-specific accumulator + /// operation. + /// + /// Each aggregation mode chooses a different `aggregate_fn` according to its + /// semantics. For example, partial aggregation takes raw inputs, and update them + /// into stored partial states, so [`GroupsAccumulator::update_batch`] is used. + pub(super) fn aggregate_batch_inner( + &mut self, + batch: &RecordBatch, + aggregate_fn: AggregateBatchFn, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + let state = self.state.building_mut(); + + let _timer = self.group_by_metrics.aggregation_time.timer(); + for group_values in &evaluated_batch.grouping_set_args { + state + .group_values + .intern(group_values, &mut state.batch_group_indices)?; + let group_indices = &state.batch_group_indices; + let total_num_groups = state.group_values.len(); + + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + aggregate_fn(acc, values, group_indices, total_num_groups)?; + } + } + + Ok(()) + } + + /// Materializes the full output once, then returns it downstream incrementally + /// by slicing it into `batch_size` chunks. + /// + /// Each aggregation mode chooses a different `materialize_accumulator_fn` + /// according to its semantics. For example, partial aggregation emits + /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`]. + /// + /// This is a temporary solution until blocked state management is implemented: + /// Issue: <https://github.com/apache/datafusion/issues/7065> + pub(super) fn next_output_batch_inner( + &mut self, + materialize_accumulator_fn: MaterializeAccumulatorFn, + ) -> Result<Option<RecordBatch>> { + let output_schema = Arc::clone(&self.output_schema); + let batch_size = self.batch_size; + + let mut output = + match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { + AggregateHashTableState::Outputting(mut state) => { + if state.group_values.is_empty() { + return Ok(None); + } + + // Accumulator output consumes internal state. Materialize all + // groups once, then slice the materialized batch on later polls. + let emit_to = EmitTo::All; + let timer = self.group_by_metrics.emitting_time.timer(); + let mut columns = state.group_values.emit(emit_to)?; + for acc in state.accumulators.iter_mut() { + columns.extend(materialize_accumulator_fn(acc, emit_to)?); + } + drop(timer); + + let batch = RecordBatch::try_new(output_schema, columns)?; + debug_assert!(batch.num_rows() > 0); + MaterializedAggregateOutput::new(batch) + } + AggregateHashTableState::OutputtingMaterialized(output) => output, + AggregateHashTableState::Done => return Ok(None), + AggregateHashTableState::Building(_) => { + return internal_err!( + "next_output_batch must be called in the outputting state" + ); + } + }; + + let batch = output.next_batch(batch_size); + if output.is_exhausted() { + self.state = AggregateHashTableState::Done; + } else { + self.state = AggregateHashTableState::OutputtingMaterialized(output); + } + Ok(batch) + } + + pub(in crate::aggregates) fn memory_size(&self) -> usize { + match &self.state { + AggregateHashTableState::Building(state) + | AggregateHashTableState::Outputting(state) => { + let acc = state + .accumulators + .iter() + .map(|acc| acc.accumulator.size()) + .sum::<usize>(); + + acc + state.group_values.size() + + state.batch_group_indices.allocated_size() + } + AggregateHashTableState::OutputtingMaterialized(output) => { + output.memory_size() + } + AggregateHashTableState::Done => 0, + } + } + + pub(in crate::aggregates) fn group_by_metrics(&self) -> &GroupByMetrics { + &self.group_by_metrics + } + + /// Returns the number of distinct groups accumulated so far. + pub(in crate::aggregates) fn building_group_count(&self) -> usize { + self.state.building().group_values.len() + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue accumulating raw input. + /// + /// Unlike normal single aggregation output, this materializes intermediate + /// states rather than final values. The states can therefore be merged after + /// spilling without finalizing the same group more than once. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + let state_schema = Arc::clone(&self.state_schema); + let state = self.state.building_mut(); + if state.group_values.is_empty() { + return Ok(None); + } + + let mut output = state.group_values.emit(EmitTo::All)?; + for acc in &mut state.accumulators { + output.extend(acc.state(EmitTo::All)?); + } + + let batch = RecordBatch::try_new(state_schema, output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is sorted for spilling. + state.group_values.clear_shrink(0); + state.batch_group_indices.clear(); + state.batch_group_indices.shrink_to_fit(); + + Ok(Some(batch)) + } + + pub(in crate::aggregates) fn is_building(&self) -> bool { + matches!(self.state, AggregateHashTableState::Building(_)) + } + + pub(in crate::aggregates) fn is_done(&self) -> bool { + matches!(self.state, AggregateHashTableState::Done) + } + + pub(super) fn start_outputting(&mut self) { + let AggregateHashTableState::Building(mut state) = + std::mem::replace(&mut self.state, AggregateHashTableState::Done) + else { + unreachable!("hash aggregate table is not building") + }; + + state.batch_group_indices = Vec::new(); + self.state = AggregateHashTableState::Outputting(state); + } +} + +/// State and argument information for a single Aggregate +/// +/// For example, for `SELECT COUNT(x), SUM(y WHERE z > 10) ...` there would be two +/// `HashAggregateAccumulator`, one each for `COUNT(x)` and `SUM(y WHERE z > 10)` +pub(super) struct HashAggregateAccumulator { + /// Aggregate expression used to create a fresh accumulator for related + /// hash tables, such as the partial-skip table. + aggregate_expr: Arc<AggregateFunctionExpr>, + + /// Arguments to pass to this accumulator. + /// + /// Example: `CORR(x, y)` stores two expressions here, while `SUM(x)` stores one. + arguments: Vec<Arc<dyn PhysicalExpr>>, + + /// Optional `FILTER` expression for this accumulator. + /// + /// Example: `SUM(x) FILTER (WHERE x > 10)` stores the `x > 10` predicate. + filter: Option<Arc<dyn PhysicalExpr>>, + + /// Accumulator state for all groups for one aggregate expression. + accumulator: Box<dyn GroupsAccumulator>, +} + +pub(super) type AggregateAccumulator = HashAggregateAccumulator; + +/// Function used by [`AggregateHashTable::aggregate_batch_inner`] to update one +/// accumulator with one evaluated input batch. +/// +/// Arguments: +/// * accumulator to update. +/// * accumulator's evaluated arguments and optional filter. +/// * one group index per input row, mapping each row to its interned group. +/// * total number of groups currently interned in that buffer, including newly +/// interned groups. +pub(super) type AggregateBatchFn = fn( + &mut AggregateAccumulator, + &EvaluatedAccumulatorArgs, + &[usize], + usize, +) -> Result<()>; + +/// Function used by [`AggregateHashTable::next_output_batch_inner`] to +/// materialize one accumulator's output columns. +/// +/// Arguments: +/// * accumulator to materialize. +/// * group range to emit from the accumulator. +pub(super) type MaterializeAccumulatorFn = + fn(&mut AggregateAccumulator, EmitTo) -> Result<Vec<ArrayRef>>; + +/// Evaluated aggregate arguments and filter for one input batch. +/// +/// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` +/// and `x > 0`. +/// +/// These arrays can be passed directly to [`GroupsAccumulator`]. +pub(super) struct EvaluatedAccumulatorArgs { + /// Evaluated argument arrays. Some aggregate functions take multiple arguments. + pub(super) arguments: Vec<ArrayRef>, + /// Evaluated filter array, `Some` if the aggregate has a `FILTER` expression. + pub(super) filter: Option<ArrayRef>, +} + +/// Evaluated all group by keys and accumulator args. +/// +/// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function evaluates +/// `k+1`, `v*v` +pub(super) struct EvaluatedAggregateBatch { + /// One entry per grouping set; each entry contains all evaluated group key + /// arrays for the current input batch. + pub(super) grouping_set_args: Vec<Vec<ArrayRef>>, + + /// Evaluated arguments and filters, one entry per aggregate expression. + pub(super) accumulator_args: Vec<EvaluatedAccumulatorArgs>, +} + +/// Buffer for the aggregate hash table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits final results during the +/// outputting stage. +/// +/// [`GroupValues`] stores the physical group-key layout, while +/// [`GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct AggregateHashTableBuffer { + /// GROUP BY expressions evaluated for each input batch. + pub(super) group_by: Arc<PhysicalGroupBy>, + + /// Interned group keys. Accumulator state is stored separately by group index. + pub(super) group_values: Box<dyn GroupValues>, + + /// Group index for each row in the current input batch. + /// + /// Each value indexes into `group_values`, and the same index is used by every + /// accumulator to update that group's aggregate state. + pub(super) batch_group_indices: Vec<usize>, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec<HashAggregateAccumulator>, +} + +pub(super) enum AggregateHashTableState { + /// Accumulating input rows into group keys and aggregate state. + Building(AggregateHashTableBuffer), + /// Emitting results directly from group keys and aggregate state. + Outputting(AggregateHashTableBuffer), + /// Materialize all the output results, and then incrementally output in the `OutputtingMaterialized` state. + /// + /// Note this is a temporary solution until the `GroupValues` issue is solved: + /// Issue: <https://github.com/apache/datafusion/issues/23178> + OutputtingMaterialized(MaterializedAggregateOutput), + Done, +} + +/// Fully evaluated aggregate output and the next row offset to emit. +/// +/// Final aggregate evaluation consumes accumulator state, and partial terminal +/// output should not repeatedly renumber group values with `EmitTo::First`. +/// Materialize once and then slice to honor `batch_size` across output polls. +pub(super) struct MaterializedAggregateOutput { + batch: RecordBatch, + offset: usize, +} + +impl MaterializedAggregateOutput { + pub(super) fn new(batch: RecordBatch) -> Self { + Self { batch, offset: 0 } + } + + pub(super) fn next_batch(&mut self, batch_size: usize) -> Option<RecordBatch> { + debug_assert!(batch_size > 0); + if self.is_exhausted() { + return None; + } + + let length = batch_size.min(self.batch.num_rows() - self.offset); + let batch = self.batch.slice(self.offset, length); + self.offset += length; + Some(batch) + } + + pub(super) fn is_exhausted(&self) -> bool { + self.offset >= self.batch.num_rows() + } + + pub(super) fn memory_size(&self) -> usize { + self.batch.get_array_memory_size() + } +} + +impl HashAggregateAccumulator { + pub(super) fn new( + aggregate_expr: Arc<AggregateFunctionExpr>, + arguments: Vec<Arc<dyn PhysicalExpr>>, + filter: Option<Arc<dyn PhysicalExpr>>, + accumulator: Box<dyn GroupsAccumulator>, + ) -> Self { + Self { + aggregate_expr, + arguments, + filter, + accumulator, + } + } + + /// Construct a new accumulator with the same definition, but with empty internal + /// state buffers (empty [`GroupsAccumulator`]). + pub(super) fn empty_like(&self) -> Result<Self> { + let accumulator = create_group_accumulator(&self.aggregate_expr)?; + Ok(Self::new( + Arc::clone(&self.aggregate_expr), + self.arguments.clone(), + self.filter.clone(), + accumulator, + )) + } + + /// Evaluate aggregate arguments and filter for one input batch. + /// + /// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` + /// and `x > 0`. + /// + /// These arrays can be passed directly to [`GroupsAccumulator`] next. + pub(super) fn evaluate_acc_args( + &self, + batch: &RecordBatch, + ) -> Result<EvaluatedAccumulatorArgs> { + let arguments = self + .arguments + .iter() + .map(|expr| { + expr.evaluate(batch) + .and_then(|value| value.into_array(batch.num_rows())) + }) + .collect::<Result<_>>()?; + + let filter = self + .filter + .as_ref() + .map(|filter| { + filter + .evaluate(batch) + .and_then(|value| value.into_array(batch.num_rows())) + }) + .transpose()?; + + Ok(EvaluatedAccumulatorArgs { arguments, filter }) + } + + pub(super) fn size(&self) -> usize { + self.accumulator.size() + } + + pub(super) fn update_batch( + &mut self, + values: &EvaluatedAccumulatorArgs, + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + let filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + self.accumulator.update_batch( + &values.arguments, + group_indices, + filter, + total_num_groups, + ) + } + + pub(super) fn merge_batch( + &mut self, + values: &EvaluatedAccumulatorArgs, + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + debug_assert!(values.filter.is_none()); + self.accumulator + .merge_batch(&values.arguments, group_indices, total_num_groups) + } + + /// Evaluating final aggregate results according to `EmitTo`, and reset inner + /// states. (e.g. after `evaluate(EmitTo::All)`, it returns all accumulated groups + /// , and clear the inner buffers) + pub(super) fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + self.accumulator.evaluate(emit_to) + } + + pub(super) fn evaluate_to_columns( + &mut self, + emit_to: EmitTo, + ) -> Result<Vec<ArrayRef>> { + Ok(vec![self.evaluate(emit_to)?]) + } + + /// Evaluating partial aggregate results according to `EmitTo`, and reset inner + /// states. (e.g. after `state(EmitTo::All)`, it returns all accumulated groups + /// , and clear the inner buffers) + pub(super) fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + self.accumulator.state(emit_to) + } + + pub(super) fn convert_to_state( + &mut self, + values: &EvaluatedAccumulatorArgs, + ) -> Result<Vec<ArrayRef>> { + let opt_filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + self.accumulator + .convert_to_state(&values.arguments, opt_filter) + } + + pub(super) fn null_arguments( + &self, + input_schema: &SchemaRef, + ) -> Result<Vec<ArrayRef>> { + self.arguments + .iter() + .map(|expr| { + let data_type = expr.data_type(input_schema)?; + Ok(new_null_array(&data_type, 1)) + }) + .collect() + } +} + +impl AggregateHashTableState { + pub(super) fn building(&self) -> &AggregateHashTableBuffer { + let Self::Building(state) = self else { + unreachable!("hash aggregate table is not building") + }; + state + } + + pub(super) fn building_mut(&mut self) -> &mut AggregateHashTableBuffer { + let Self::Building(state) = self else { + unreachable!("hash aggregate table is not building") + }; + state + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Array, Int32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + + use super::*; + + #[test] + fn materialized_aggregate_output_slices_batches_until_exhausted() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group_col", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + )?; + let mut output = MaterializedAggregateOutput::new(batch); + + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]); + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]); + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]); + assert!(output.next_batch(2).is_none()); + assert!(output.is_exhausted()); + + Ok(()) + } + + fn int32_values(batch: &RecordBatch, column: usize) -> Vec<i32> { + let array = batch + .column(column) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + (0..array.len()).map(|idx| array.value(idx)).collect() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs new file mode 100644 index 0000000..2293e7b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
@@ -0,0 +1,410 @@ +// 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. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::assert_or_internal_err; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; + +/// Aggregate table shared by the ordered partial and final paths. +/// +/// # Ordering optimization +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// are proven complete. Completed groups can be emitted before the input stream +/// ends, which keeps memory bounded by the active ordered key range. +/// +/// # Partial and final variant difference +/// +/// The partial and final aggregate tables implement the two stages of grouped +/// aggregation. See +/// [`OrderedPartialAggregateStream`](crate::aggregates::ordered_partial_stream::OrderedPartialAggregateStream) +/// for the high-level plan shape. +/// +/// Example: `AVG(v) FILTER (WHERE v>0) GROUP BY k` +/// +/// Partial table ([`AggregateMode::Partial`], with optional filter from query): +/// - Input rows: `k, v` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, sum(v), count(v)` +/// +/// Final table ([`AggregateMode::Final`], no filters): +/// - Input rows: `k, sum(v), count(v)` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, avg(v)` +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::<PartialMarker>::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::<FinalMarker>::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl<T>`; partial/final behavior lives on +/// marker-specific impls. +pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> { + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Intermediate-state schema used when memory pressure requires the table + /// to pass through or spill its current state. + pub(super) state_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData<OrderedAggrMode>, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the +/// input ordering proves those groups are complete. +/// +/// [`GroupOrdering`] tracks when and how to do early emit. +/// [`GroupValues`] stores the physical group-key layout, while +/// [`datafusion_expr::GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct OrderedAggregateTableBuffer { + /// GROUP BY expressions evaluated against input batches. + pub(super) group_by: Arc<PhysicalGroupBy>, + + /// Tracks how far ordered input allows this table to drain safely. + pub(super) group_ordering: GroupOrdering, + + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box<dyn GroupValues>, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec<usize>, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec<AggregateAccumulator>, +} + +/// Methods shared by all aggregate modes +impl<AggrMode> OrderedAggregateTable<AggrMode> { + #[expect( + clippy::too_many_arguments, + reason = "keeps ordered partial and final table construction explicit" + )] + pub(super) fn new_for_mode( + agg: &AggregateExec, + input_schema: &SchemaRef, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + aggregate_mode: &AggregateMode, + filters: Vec<Option<Arc<dyn PhysicalExpr>>>, + group_by_metrics: GroupByMetrics, + ) -> Result<Self> { + assert_or_internal_err!( + batch_size > 0, + "OrderedAggregateTable requires config batch_size >= 1" + ); + + let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_schema = agg.group_by.group_schema(input_schema)?; + let group_values = new_group_values(group_schema, &group_ordering)?; + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + aggregate_mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(AggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::<Result<_>>()?; + + Ok(Self { + output_schema, + state_schema, + batch_size, + group_by_metrics, + buffer: OrderedAggregateTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_ordering, + group_values, + group_indices: vec![], + accumulators, + }, + _mode: PhantomData, + }) + } + + /// Evaluates all group by keys and accumulator args. + /// + /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function + /// evaluates `k+1`, `v*v`. + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result<EvaluatedAggregateBatch> { + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + let grouping_set_args = evaluate_group_by(&self.buffer.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + let accumulator_args = self + .buffer + .accumulators + .iter() + .map(|acc| acc.evaluate_acc_args(batch)) + .collect::<Result<Vec<_>>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Called after the input stream is exhausted and the last batch has been + /// aggregated. + /// + /// Updates the internal `GroupOrdering` so it can continue emitting until + /// the buffer is empty. + pub(in crate::aggregates) fn input_done(&mut self) { + self.buffer.group_ordering.input_done(); + } + + /// Returns the ordering state used to decide how memory pressure is handled. + pub(in crate::aggregates) fn group_ordering(&self) -> &GroupOrdering { + &self.buffer.group_ordering + } + + /// Number of groups currently buffered. + pub(in crate::aggregates) fn num_groups(&self) -> usize { + self.buffer.group_values.len() + } + + /// Check if there is zero groups accumulated so far. + pub(in crate::aggregates) fn is_empty(&self) -> bool { + self.num_groups() == 0 + } + + /// All internal buffer's memory size. + pub(in crate::aggregates) fn memory_size(&self) -> usize { + self.buffer + .accumulators + .iter() + .map(|acc| acc.size()) + .sum::<usize>() + + self.buffer.group_values.size() + + self.buffer.group_ordering.size() + + self.buffer.group_indices.allocated_size() + } + + pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics { + self.group_by_metrics.clone() + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue with a new ordered input segment. + /// + /// Unlike normal ordered emission, this operation is allowed to take the + /// active (incomplete) groups. Partial aggregation can pass those states to + /// its final stage, while final aggregation sorts and spills them before + /// replay. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let mut output = self.buffer.group_values.emit(EmitTo::All)?; + for acc in &mut self.buffer.accumulators { + output.extend(acc.state(EmitTo::All)?); + } + + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is passed downstream or sorted for spilling. + self.buffer.group_values.clear_shrink(0); + self.buffer.group_indices.clear(); + self.buffer.group_indices.shrink_to_fit(); + self.buffer.group_ordering.reset(); + + Ok(Some(batch)) + } + + /// Returns the [`EmitTo`], clamped to the specified batch size + /// + /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number + /// of groups to emit from `GroupValues` / accumulators, and + /// `should_remove_groups` indicates whether `GroupOrdering` must also shift + /// its tracked indexes. + pub(super) fn clamp_emit_to( + &self, + group_count: usize, + emit_to: EmitTo, + ) -> (EmitTo, bool) { + match emit_to { + EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), + EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), + EmitTo::All => (EmitTo::First(self.batch_size), false), + } + } + /// Aggregates one evaluated input batch. + /// + /// This common utility is used by ordered partial and ordered final aggregation. + /// + /// # Argument: `is_final` + /// + /// - `true`: merge partial aggregate states for final aggregation. + /// - `false`: update aggregate states from raw input for partial aggregation. + pub(super) fn aggregate_evaluated_batch( + &mut self, + evaluated_batch: &EvaluatedAggregateBatch, + is_final: bool, + ) -> Result<()> { + for group_values in &evaluated_batch.grouping_set_args { + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); + if total_num_groups > starting_num_groups { + self.buffer.group_ordering.new_groups( + group_values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + + let timer = self.group_by_metrics.aggregation_time.timer(); + for (acc, values) in self + .buffer + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + if is_final { + acc.merge_batch( + values, + &self.buffer.group_indices, + total_num_groups, + )?; + } else { + acc.update_batch( + values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + } + drop(timer); + } + + Ok(()) + } + + /// Emits groups allowed by `GroupOrdering`, leaving only the current + /// unfinished ordered-key range buffered. + /// + /// This common utility is used by ordered partial and ordered final aggregation. + /// + /// # Argument: `is_final` + /// + /// - `true`: output final aggregate values. + /// - `false`: output partial accumulator states. + pub(super) fn next_output_batch_for_mode( + &mut self, + is_final: bool, + ) -> Result<Option<RecordBatch>> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let Some(emit_to) = self.buffer.group_ordering.emit_to() else { + return Ok(None); + }; + let (emit_to, should_remove_groups) = + self.clamp_emit_to(self.buffer.group_values.len(), emit_to); + + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.buffer.group_values.emit(emit_to)?; + if should_remove_groups { + match emit_to { + EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), + // `EmitTo::All` is only used after `input_done`, when all + // buffered groups are known complete and the ordering state is + // no longer needed. + EmitTo::All => {} + } + } + + for acc in &mut self.buffer.accumulators { + if is_final { + output.push(acc.evaluate(emit_to)?); + } else { + output.extend(acc.state(emit_to)?); + } + } + drop(timer); + + let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + Ok(Some(batch)) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs new file mode 100644 index 0000000..b80e15d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
@@ -0,0 +1,77 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; + +use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +impl AggregateHashTable<FinalMarker> { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result<Self> { + Self::new_with_filters( + agg, + partition, + output_schema, + Arc::clone(&agg.input().schema()), + batch_size, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Emits the next batch of aggregated group keys and final aggregate values. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) + } + + /// Final aggregation consumes partial aggregate states and merges them into + /// the table's partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/mod.rs new file mode 100644 index 0000000..2c7ec01 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/mod.rs
@@ -0,0 +1,31 @@ +// 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. + +mod common; +mod common_ordered; +mod final_table; +mod ordered_final_table; +mod ordered_partial_table; +mod partial_reduce_table; +mod partial_table; +mod single_table; + +pub(super) use common::{ + AggregateHashTable, FinalMarker, PartialMarker, PartialReduceMarker, + PartialSkipMarker, SingleMarker, +}; +pub(super) use common_ordered::OrderedAggregateTable;
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs new file mode 100644 index 0000000..fd064eb --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
@@ -0,0 +1,85 @@ +// 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. + +//! Aggregate table for final aggregation when partial-state input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::InputOrderMode; +use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::group_values::GroupByMetrics; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable<FinalMarker> { + pub(in crate::aggregates) fn new_with_input_order( + agg: &AggregateExec, + input_schema: &SchemaRef, + output_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + group_by_metrics: GroupByMetrics, + ) -> Result<Self> { + Self::new_for_mode( + agg, + input_schema, + output_schema, + Arc::clone(input_schema), + batch_size, + input_order_mode, + &AggregateMode::Final, + vec![None; agg.aggr_expr.len()], + group_by_metrics, + ) + } + + /// Merges one partial-state input batch and updates ordering information for + /// any newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + // `PhysicalGroupBy::as_final()` removes grouping sets while planning + // final aggregation, so final ordered aggregation sees one grouping. + debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1); + self.aggregate_evaluated_batch(&evaluated_batch, true) + } + + /// See comments in `ordered_partial_stream::next_output_batch` + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + self.next_output_batch_for_mode(true) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs new file mode 100644 index 0000000..a04e4dd --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
@@ -0,0 +1,106 @@ +// 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. + +//! Aggregate table for partial aggregation when input is ordered by group keys. +//! +//! See the [`super::common_ordered`] comments for the high-level ideas. +//! +//! This operator handles input that is ordered by group keys: +//! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b` +//! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a` +//! +//! When a group key combination is exhausted, this table eagerly flushes the +//! completed groups to improve memory efficiency. +//! +//! The implementation is separated from other aggregate tables because this +//! execution path is likely to be optimized further in the future. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::{ + AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker, + group_values::GroupByMetrics, +}; + +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable<PartialMarker> { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result<Self> { + let input_schema = agg.input().schema(); + let state_schema = Arc::clone(&output_schema); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + Self::new_for_mode( + agg, + &input_schema, + output_schema, + state_schema, + batch_size, + &agg.input_order_mode, + &AggregateMode::Partial, + agg.filter_expr.iter().cloned().collect(), + group_by_metrics, + ) + } + + /// Aggregates one raw input batch and updates ordering information for any + /// newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + self.aggregate_evaluated_batch(&evaluated_batch, false) + } + + /// Emits the next batch of partial state rows for groups proven complete by + /// the input ordering. + /// + /// For example, when the query is `GROUP BY a` and the input is ordered by + /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` + /// are complete and safe to emit. + /// + /// Key steps: + /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly. + /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and + /// all `GroupsAccumulator`s. + /// + /// This may output small batches. Avoiding tiny batches is left to future + /// ordered-aggregation optimizations. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + self.next_output_batch_for_mode(false) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs new file mode 100644 index 0000000..4dfd6a7 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
@@ -0,0 +1,71 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; + +use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker}; + +/// Methods specific to the aggregate hash table used in the partial-reduce stage. +impl AggregateHashTable<PartialReduceMarker> { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result<Self> { + Self::new_with_filters( + agg, + partition, + Arc::clone(&output_schema), + output_schema, + batch_size, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Emits the next batch of aggregated group keys and aggregate states. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + self.next_output_batch_inner(HashAggregateAccumulator::state) + } + + /// Partial-reduce aggregation consumes partial aggregate states and merges + /// them into the table's partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs new file mode 100644 index 0000000..a64fd32 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs
@@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, new_null_array}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; + +use crate::aggregates::group_values::new_group_values; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal}; + +use super::common::{ + AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, + EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, +}; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` +impl AggregateHashTable<PartialMarker> { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result<Self> { + Self::new_with_filters( + agg, + partition, + Arc::clone(&output_schema), + output_schema, + batch_size, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// Emits the next batch of aggregated group keys and aggregate states. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + self.next_output_batch_inner(HashAggregateAccumulator::state) + } + + /// In skip-partial-aggregation optimization, when a decision has been made to skip + /// partial stage, build a typed hash table only for aggregation state conversion + /// row-by-row. + pub(in crate::aggregates) fn partial_skip_table( + &self, + ) -> Result<AggregateHashTable<PartialSkipMarker>> { + let state = self.state.building(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let accumulators = state + .accumulators + .iter() + .map(HashAggregateAccumulator::empty_like) + .collect::<Result<Vec<_>>>()?; + + Ok(AggregateHashTable { + group_by_metrics: self.group_by_metrics.clone(), + input_schema: Arc::clone(&self.input_schema), + output_schema: Arc::clone(&self.output_schema), + state_schema: Arc::clone(&self.state_schema), + batch_size: self.batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&state.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + _mode: PhantomData, + }) + } + + /// Partial aggregation consumes raw input rows and updates the table's + /// partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.init_empty_grouping_sets()?; + self.start_outputting(); + Ok(()) + } + + /// Creates the required empty grouping-set rows when the input is empty. + /// + /// For example, this query must still produce one grand-total group even if + /// `t` has no rows: + /// + /// ```sql + /// SELECT COUNT(v) + /// FROM t + /// GROUP BY GROUPING SETS (()); + /// ``` + /// + /// The synthetic row is filtered out before accumulator update so aggregates + /// see the same state they would see for an empty input, rather than a real + /// null-valued row. + fn init_empty_grouping_sets(&mut self) -> Result<()> { + let state = self.state.building_mut(); + if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + return Ok(()); + } + + let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let n_expr = state.group_by.expr().len(); + let mut any_interned = false; + + for group in state.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let ordinal = *entry; + *entry += 1; + ordinal + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + let mut cols: Vec<ArrayRef> = group_schema + .fields() + .iter() + .take(n_expr) + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + state + .group_values + .intern(&cols, &mut state.batch_group_indices)?; + any_interned = true; + } + + if any_interned { + let total_groups = state.group_values.len(); + let false_filter = BooleanArray::from(vec![false]); + for acc in state.accumulators.iter_mut() { + let null_args = acc.null_arguments(&self.input_schema)?; + let values = EvaluatedAccumulatorArgs { + arguments: null_args, + filter: Some(Arc::new(false_filter.clone())), + }; + acc.update_batch(&values, &[0], total_groups)?; + } + } + + Ok(()) + } +} + +impl AggregateHashTable<PartialSkipMarker> { + pub(in crate::aggregates) fn convert_batch_to_state( + &mut self, + batch: &RecordBatch, + ) -> Result<RecordBatch> { + let evaluated_batch = self.evaluate_batch(batch)?; + + assert_eq_or_internal_err!( + evaluated_batch.grouping_set_args.len(), + 1, + "group_values expected to have single element" + ); + let mut output = evaluated_batch + .grouping_set_args + .into_iter() + .next() + .unwrap_or_default(); + + let state = self.state.building_mut(); + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + output.extend(acc.convert_to_state(values)?); + } + + Ok(RecordBatch::try_new( + Arc::clone(&self.output_schema), + output, + )?) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs new file mode 100644 index 0000000..56d601c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
@@ -0,0 +1,76 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; + +use super::common::{AggregateHashTable, HashAggregateAccumulator, SingleMarker}; + +/// Implementation specific to single aggregation, where the table stores final +/// aggregate values and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, avg(x)` +/// - Input rows: `k, x` +impl AggregateHashTable<SingleMarker> { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + ) -> Result<Self> { + Self::new_with_filters( + agg, + partition, + output_schema, + state_schema, + batch_size, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// Emits the next batch of aggregated group keys and final aggregate values. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result<Option<RecordBatch>> { + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) + } + + /// Single aggregation consumes raw input rows and updates the table's + /// final-value accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_stream.rs new file mode 100644 index 0000000..ac7727b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/aggregate_stream.rs
@@ -0,0 +1,478 @@ +// 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. + +//! Aggregate without grouping columns + +use crate::aggregates::{ + AccumulatorItem, AggrDynFilter, AggregateInputMode, AggregateMode, + DynamicFilterAggregateType, aggregate_expressions, create_accumulators, + finalize_aggregation, +}; +use crate::metrics::{BaselineMetrics, RecordOutput}; +use crate::stream::EmptyRecordBatchStream; +use crate::{RecordBatchStream, SendableRecordBatchStream}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, ScalarValue, internal_datafusion_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{BinaryExpr, lit}; +use futures::stream::BoxStream; +use std::borrow::Cow; +use std::cmp::Ordering; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::AggregateExec; +use crate::filter::batch_filter; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::stream::{Stream, StreamExt}; + +/// stream struct for aggregation without grouping columns +pub(crate) struct AggregateStream { + stream: BoxStream<'static, Result<RecordBatch>>, + schema: SchemaRef, +} + +/// Actual implementation of [`AggregateStream`]. +/// +/// This is wrapped into yet another struct because we need to interact with the async memory management subsystem +/// during poll. To have as little code "weirdness" as possible, we chose to just use [`BoxStream`] together with +/// [`futures::stream::unfold`]. +/// +/// The latter requires a state object, which is [`AggregateStreamInner`]. +struct AggregateStreamInner { + // ==== Properties ==== + schema: SchemaRef, + mode: AggregateMode, + input: SendableRecordBatchStream, + aggregate_expressions: Vec<Vec<Arc<dyn PhysicalExpr>>>, + filter_expressions: Arc<[Option<Arc<dyn PhysicalExpr>>]>, + + // ==== Runtime States/Buffers ==== + accumulators: Vec<AccumulatorItem>, + // None if the dynamic filter is not applicable. See details in `AggrDynFilter`. + agg_dyn_filter_state: Option<Arc<AggrDynFilter>>, + finished: bool, + + // ==== Execution Resources ==== + baseline_metrics: BaselineMetrics, + reservation: MemoryReservation, +} + +impl AggregateStreamInner { + // TODO: check if we get Null handling correct + /// # Examples + /// - Example 1 + /// Accumulators: min(c1) + /// Current Bounds: min(c1)=10 + /// --> dynamic filter PhysicalExpr: c1 < 10 + /// + /// - Example 2 + /// Accumulators: min(c1), max(c1), min(c2) + /// Current Bounds: min(c1)=10, max(c1)=100, min(c2)=20 + /// --> dynamic filter PhysicalExpr: (c1 < 10) OR (c1>100) OR (c2 < 20) + /// + /// # Errors + /// Returns internal errors if the dynamic filter is not enabled, or other + /// invariant check fails. + fn build_dynamic_filter_from_accumulator_bounds( + &self, + ) -> Result<Arc<dyn PhysicalExpr>> { + let Some(filter_state) = self.agg_dyn_filter_state.as_ref() else { + return internal_err!( + "`build_dynamic_filter_from_accumulator_bounds()` is only called when dynamic filter is enabled" + ); + }; + + let mut predicates: Vec<Arc<dyn PhysicalExpr>> = + Vec::with_capacity(filter_state.supported_accumulators_info.len()); + + for acc_info in &filter_state.supported_accumulators_info { + // Skip if we don't yet have a meaningful bound + let bound = { + let guard = acc_info.shared_bound.lock(); + if (*guard).is_null() { + continue; + } + guard.clone() + }; + + let agg_exprs = self + .aggregate_expressions + .get(acc_info.aggr_index) + .ok_or_else(|| { + internal_datafusion_err!( + "Invalid aggregate expression index {} for dynamic filter", + acc_info.aggr_index + ) + })?; + // Only aggregates with a single argument are supported. + let column_expr = agg_exprs.first().ok_or_else(|| { + internal_datafusion_err!( + "Aggregate expression at index {} expected a single argument", + acc_info.aggr_index + ) + })?; + + let literal = lit(bound); + let predicate: Arc<dyn PhysicalExpr> = match acc_info.aggr_type { + DynamicFilterAggregateType::Min => Arc::new(BinaryExpr::new( + Arc::clone(column_expr), + Operator::Lt, + literal, + )), + DynamicFilterAggregateType::Max => Arc::new(BinaryExpr::new( + Arc::clone(column_expr), + Operator::Gt, + literal, + )), + }; + predicates.push(predicate); + } + + let combined = predicates.into_iter().reduce(|acc, pred| { + Arc::new(BinaryExpr::new(acc, Operator::Or, pred)) as Arc<dyn PhysicalExpr> + }); + + Ok(combined.unwrap_or_else(|| lit(true))) + } + + // If the dynamic filter is enabled, update it using the current accumulator's + // values + fn maybe_update_dyn_filter(&mut self) -> Result<()> { + // Step 1: Update each partition's current bound + let Some(filter_state) = self.agg_dyn_filter_state.as_ref() else { + return Ok(()); + }; + + let mut bounds_changed = false; + + for acc_info in &filter_state.supported_accumulators_info { + let acc = + self.accumulators + .get_mut(acc_info.aggr_index) + .ok_or_else(|| { + internal_datafusion_err!( + "Invalid accumulator index {} for dynamic filter", + acc_info.aggr_index + ) + })?; + // First get current partition's bound, then update the shared bound among + // all partitions. + let current_bound = acc.evaluate()?; + { + let mut bound = acc_info.shared_bound.lock(); + let new_bound = match acc_info.aggr_type { + DynamicFilterAggregateType::Max => { + scalar_max(&bound, ¤t_bound)? + } + DynamicFilterAggregateType::Min => { + scalar_min(&bound, ¤t_bound)? + } + }; + if new_bound != *bound { + *bound = new_bound; + bounds_changed = true; + } + } + } + + // Step 2: Sync the dynamic filter physical expression with reader, + // but only if any bound actually changed. + if bounds_changed { + let predicate = self.build_dynamic_filter_from_accumulator_bounds()?; + filter_state.filter.update(predicate)?; + } + + Ok(()) + } +} + +/// Returns the element-wise minimum of two `ScalarValue`s. +/// +/// # Null semantics +/// - `min(NULL, NULL) = NULL` +/// - `min(NULL, x) = x` +/// - `min(x, NULL) = x` +/// +/// # Errors +/// Returns internal error if v1 and v2 has incompatible types. +fn scalar_min(v1: &ScalarValue, v2: &ScalarValue) -> Result<ScalarValue> { + if let Some(result) = scalar_cmp_null_short_circuit(v1, v2) { + return Ok(result); + } + + match v1.partial_cmp(v2) { + Some(Ordering::Less | Ordering::Equal) => Ok(v1.clone()), + Some(Ordering::Greater) => Ok(v2.clone()), + None => datafusion_common::internal_err!( + "cannot compare values of different or incompatible types: {v1:?} vs {v2:?}" + ), + } +} + +/// Returns the element-wise maximum of two `ScalarValue`s. +/// +/// # Null semantics +/// - `max(NULL, NULL) = NULL` +/// - `max(NULL, x) = x` +/// - `max(x, NULL) = x` +/// +/// # Errors +/// Returns internal error if v1 and v2 has incompatible types. +fn scalar_max(v1: &ScalarValue, v2: &ScalarValue) -> Result<ScalarValue> { + if let Some(result) = scalar_cmp_null_short_circuit(v1, v2) { + return Ok(result); + } + + match v1.partial_cmp(v2) { + Some(Ordering::Greater | Ordering::Equal) => Ok(v1.clone()), + Some(Ordering::Less) => Ok(v2.clone()), + None => datafusion_common::internal_err!( + "cannot compare values of different or incompatible types: {v1:?} vs {v2:?}" + ), + } +} + +fn scalar_cmp_null_short_circuit( + v1: &ScalarValue, + v2: &ScalarValue, +) -> Option<ScalarValue> { + match (v1, v2) { + (ScalarValue::Null, ScalarValue::Null) => Some(ScalarValue::Null), + (ScalarValue::Null, other) | (other, ScalarValue::Null) => Some(other.clone()), + _ => None, + } +} + +/// Prepend the grouping ID column to the output columns if present. +/// +/// For GROUPING SETS with no GROUP BY expressions, the schema includes a `__grouping_id` +/// column that must be present in the output. This function inserts it at the beginning +/// of the columns array to maintain schema alignment. +fn prepend_grouping_id_column( + mut columns: Vec<Arc<dyn arrow::array::Array>>, + grouping_id: Option<&ScalarValue>, +) -> Result<Vec<Arc<dyn arrow::array::Array>>> { + if let Some(id) = grouping_id { + let num_rows = columns.first().map(|array| array.len()).unwrap_or(1); + let grouping_ids = id.to_array_of_size(num_rows)?; + columns.insert(0, grouping_ids); + } + Ok(columns) +} + +impl AggregateStream { + /// Create a new AggregateStream + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + let agg_schema = Arc::clone(&agg.schema); + let agg_filter_expr = Arc::clone(&agg.filter_expr); + + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let input = agg.input.execute(partition, Arc::clone(context))?; + + let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?; + let filter_expressions = match agg.mode.input_mode() { + AggregateInputMode::Raw => agg_filter_expr, + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), + }; + let accumulators = create_accumulators(&agg.aggr_expr)?; + + let reservation = MemoryConsumer::new(format!("AggregateStream[{partition}]")) + .register(context.memory_pool()); + + // Enable dynamic filter if: + // 1. AggregateExec did the check and ensure it supports the dynamic filter + // (its dynamic_filter field will be Some(..)) + // 2. Aggregate dynamic filter is enabled from the config + let mut maybe_dynamic_filter = match agg.dynamic_filter.as_ref() { + Some(filter) => Some(Arc::clone(filter)), + _ => None, + }; + + if !context + .session_config() + .options() + .optimizer + .enable_aggregate_dynamic_filter_pushdown + { + maybe_dynamic_filter = None; + } + + let inner = AggregateStreamInner { + schema: Arc::clone(&agg.schema), + mode: agg.mode, + input, + baseline_metrics, + aggregate_expressions, + filter_expressions, + accumulators, + reservation, + finished: false, + agg_dyn_filter_state: maybe_dynamic_filter, + }; + + let stream = futures::stream::unfold(inner, |mut this| async move { + if this.finished { + return None; + } + + loop { + let result = match this.input.next().await { + Some(Ok(batch)) => { + let result = { + let elapsed_compute = this.baseline_metrics.elapsed_compute(); + let _timer = elapsed_compute.timer(); // Stops on drop + aggregate_batch( + &this.mode, + &batch, + &mut this.accumulators, + &this.aggregate_expressions, + &this.filter_expressions, + ) + }; + + let result = result.and_then(|allocated| { + this.maybe_update_dyn_filter()?; + Ok(allocated) + }); + + // allocate memory + // This happens AFTER we actually used the memory, but simplifies the whole accounting and we are OK with + // overshooting a bit. Also this means we either store the whole record batch or not. + match result + .and_then(|allocated| this.reservation.try_grow(allocated)) + { + Ok(_) => continue, + Err(e) => Err(e), + } + } + Some(Err(e)) => Err(e), + None => { + this.finished = true; + // Release the input pipeline's resources before finalization. + let input_schema = this.input.schema(); + this.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + let timer = this.baseline_metrics.elapsed_compute().timer(); + let result = + finalize_aggregation(&mut this.accumulators, &this.mode) + .and_then(|columns| { + prepend_grouping_id_column(columns, None) + }) + .and_then(|columns| { + RecordBatch::try_new( + Arc::clone(&this.schema), + columns, + ) + .map_err(Into::into) + }) + .record_output(&this.baseline_metrics); + + timer.done(); + + result + } + }; + + this.finished = true; + return Some((result, this)); + } + }); + + // seems like some consumers call this stream even after it returned `None`, so let's fuse the stream. + let stream = stream.fuse(); + let stream = Box::pin(stream); + + Ok(Self { + schema: agg_schema, + stream, + }) + } +} + +impl Stream for AggregateStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let this = &mut *self; + this.stream.poll_next_unpin(cx) + } +} + +impl RecordBatchStream for AggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Perform group-by aggregation for the given [`RecordBatch`]. +/// +/// If successful, this returns the additional number of bytes that were allocated during this process. +/// +/// TODO: Make this a member function +fn aggregate_batch( + mode: &AggregateMode, + batch: &RecordBatch, + accumulators: &mut [AccumulatorItem], + expressions: &[Vec<Arc<dyn PhysicalExpr>>], + filters: &[Option<Arc<dyn PhysicalExpr>>], +) -> Result<usize> { + let mut allocated = 0usize; + + // 1.1 iterate accumulators and respective expressions together + // 1.2 filter the batch if necessary + // 1.3 evaluate expressions + // 1.4 update / merge accumulators with the expressions' values + + // 1.1 + accumulators + .iter_mut() + .zip(expressions) + .zip(filters) + .try_for_each(|((accum, expr), filter)| { + // 1.2 + let batch = match filter { + Some(filter) => Cow::Owned(batch_filter(batch, filter)?), + None => Cow::Borrowed(batch), + }; + + // 1.3 + let values = evaluate_expressions_to_arrays(expr, batch.as_ref())?; + + // 1.4 + let size_pre = accum.size(); + let res = match mode.input_mode() { + AggregateInputMode::Raw => accum.update_batch(&values), + AggregateInputMode::Partial => accum.merge_batch(&values), + }; + let size_post = accum.size(); + allocated += size_post.saturating_sub(size_pre); + res + })?; + + Ok(allocated) +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/metrics.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/metrics.rs new file mode 100644 index 0000000..1c6285d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/metrics.rs
@@ -0,0 +1,222 @@ +// 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. + +//! Metrics for the various group-by implementations. + +use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; + +#[derive(Clone)] +pub(crate) struct GroupByMetrics { + /// Time spent calculating the group IDs from the evaluated grouping columns. + pub(crate) time_calculating_group_ids: Time, + /// Time spent evaluating the inputs to the aggregate functions. + pub(crate) aggregate_arguments_time: Time, + /// Time spent evaluating the aggregate expressions themselves + /// (e.g. summing all elements and counting number of elements for `avg` aggregate). + pub(crate) aggregation_time: Time, + /// Time spent emitting the final results and constructing the record batch + /// which includes finalizing the grouping expressions + /// (e.g. emit from the hash table in case of hash aggregation) and the accumulators + pub(crate) emitting_time: Time, +} + +impl GroupByMetrics { + pub(crate) fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + time_calculating_group_ids: MetricBuilder::new(metrics) + .subset_time("time_calculating_group_ids", partition), + aggregate_arguments_time: MetricBuilder::new(metrics) + .subset_time("aggregate_arguments_time", partition), + aggregation_time: MetricBuilder::new(metrics) + .subset_time("aggregation_time", partition), + emitting_time: MetricBuilder::new(metrics) + .subset_time("emitting_time", partition), + } + } +} + +#[cfg(test)] +mod tests { + use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use crate::metrics::MetricsSet; + use crate::test::TestMemoryExec; + use crate::{ExecutionPlan, collect}; + use arrow::array::{Float64Array, UInt32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion_common::Result; + use datafusion_execution::TaskContext; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use std::sync::Arc; + + /// Helper function to verify all three GroupBy metrics exist and have non-zero values + fn assert_groupby_metrics(metrics: &MetricsSet) { + let agg_arguments_time = metrics.sum_by_name("aggregate_arguments_time"); + assert!(agg_arguments_time.is_some()); + assert!(agg_arguments_time.unwrap().as_usize() > 0); + + let aggregation_time = metrics.sum_by_name("aggregation_time"); + assert!(aggregation_time.is_some()); + assert!(aggregation_time.unwrap().as_usize() > 0); + + let emitting_time = metrics.sum_by_name("emitting_time"); + assert!(emitting_time.is_some()); + assert!(emitting_time.unwrap().as_usize() > 0); + } + + #[tokio::test] + async fn test_groupby_metrics_partial_mode() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + // Create multiple batches to ensure metrics accumulate + let batches = (0..5) + .map(|i| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3, 4])), + Arc::new(Float64Array::from(vec![ + i as f64, + (i + 1) as f64, + (i + 2) as f64, + (i + 3) as f64, + ])), + ], + ) + .unwrap() + }) + .collect::<Vec<_>>(); + + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let aggregates = vec![ + Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(b)") + .build()?, + ), + ]; + + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggregates, + vec![None, None], + input, + schema, + )?); + + // This test is for `GroupByMetrics`, which are maintained by + // `GroupedHashAggregateStream`. Use a finite memory pool so the partial + // aggregate does not take the initial-partial stream path. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(10 * 1024 * 1024, 1.0) + .build_arc()?; + let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); + let _result = + collect(Arc::clone(&aggregate_exec) as _, Arc::clone(&task_ctx)).await?; + + let metrics = aggregate_exec.metrics().unwrap(); + assert_groupby_metrics(&metrics); + + Ok(()) + } + + #[tokio::test] + async fn test_groupby_metrics_final_mode() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + let batches = (0..3) + .map(|i| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3])), + Arc::new(Float64Array::from(vec![ + i as f64, + (i + 1) as f64, + (i + 2) as f64, + ])), + ], + ) + .unwrap() + }) + .collect::<Vec<_>>(); + + let partial_input = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + // Create partial aggregate + let partial_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + partial_input, + Arc::clone(&schema), + )?); + + // Create final aggregate + let final_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggregates, + vec![None], + partial_aggregate, + schema, + )?); + + let task_ctx = Arc::new(TaskContext::default()); + let _result = + collect(Arc::clone(&final_aggregate) as _, Arc::clone(&task_ctx)).await?; + + let metrics = final_aggregate.metrics().unwrap(); + assert_groupby_metrics(&metrics); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs new file mode 100644 index 0000000..1101d53 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/mod.rs
@@ -0,0 +1,214 @@ +// 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. + +//! [`GroupValues`] trait for storing and interning group keys + +use arrow::array::types::{ + Date32Type, Date64Type, Decimal128Type, Time32MillisecondType, Time32SecondType, + Time64MicrosecondType, Time64NanosecondType, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, +}; +use arrow::array::{ArrayRef, downcast_primitive}; +use arrow::datatypes::{DataType, SchemaRef, TimeUnit}; +use datafusion_common::Result; + +use datafusion_expr::EmitTo; + +pub mod multi_group_by; + +mod row; +pub use row::GroupValuesRows; +mod single_group_by; +use datafusion_physical_expr::binary_map::OutputType; +use multi_group_by::GroupValuesColumn; + +pub(crate) use single_group_by::primitive::HashValue; + +use crate::aggregates::{ + group_values::single_group_by::{ + boolean::GroupValuesBoolean, bytes::GroupValuesBytes, + bytes_view::GroupValuesBytesView, primitive::GroupValuesPrimitive, + }, + order::GroupOrdering, +}; + +mod metrics; +mod null_builder; + +pub(crate) use metrics::GroupByMetrics; + +/// Stores the group values during hash aggregation. +/// +/// # Background +/// +/// In a query such as `SELECT a, b, count(*) FROM t GROUP BY a, b`, the group values +/// identify each group, and correspond to all the distinct values of `(a,b)`. +/// +/// ```sql +/// -- Input has 4 rows with 3 distinct combinations of (a,b) ("groups") +/// create table t(a int, b varchar) +/// as values (1, 'a'), (2, 'b'), (1, 'a'), (3, 'c'); +/// +/// select a, b, count(*) from t group by a, b; +/// ---- +/// 1 a 2 +/// 2 b 1 +/// 3 c 1 +/// ``` +/// +/// # Design +/// +/// Managing group values is a performance critical operation in hash +/// aggregation. The major operations are: +/// +/// 1. Intern: Quickly finding existing and adding new group values +/// 2. Emit: Returning the group values as an array +/// +/// There are multiple specialized implementations of this trait optimized for +/// different data types and number of columns, optimized for these operations. +/// See [`new_group_values`] for details. +/// +/// # Group Ids +/// +/// Each distinct group in a hash aggregation is identified by a unique group id +/// (usize) which is assigned by instances of this trait. Group ids are +/// continuous without gaps, starting from 0. +pub trait GroupValues: Send { + /// Calculates the group id for each input row of `cols`, assigning new + /// group ids as necessary. + /// + /// When the function returns, `groups` must contain the group id for each + /// row in `cols`. + /// + /// If a row has the same value as a previous row, the same group id is + /// assigned. If a row has a new value, the next available group id is + /// assigned. + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()>; + + /// Returns the number of bytes of memory used by this [`GroupValues`]. + /// + /// May be expensive; check the implementation before calling on hot paths. + fn size(&self) -> usize; + + /// Returns true if this [`GroupValues`] is empty + fn is_empty(&self) -> bool; + + /// The number of values (distinct group values) stored in this [`GroupValues`] + fn len(&self) -> usize; + + /// Emits the group values + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>; + + /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage) + fn clear_shrink(&mut self, num_rows: usize); +} + +/// Return a specialized implementation of [`GroupValues`] for the given schema. +/// +/// [`GroupValues`] implementations choosing logic: +/// +/// - If group by single column, and type of this column has +/// the specific [`GroupValues`] implementation, such implementation +/// will be chosen. +/// +/// - If group by multiple columns, and all column types have the specific +/// `GroupColumn` implementations, `GroupValuesColumn` will be chosen. +/// +/// - Otherwise, the general implementation `GroupValuesRows` will be chosen. +/// +/// `GroupColumn`: crate::aggregates::group_values::multi_group_by::GroupColumn +/// `GroupValuesColumn`: crate::aggregates::group_values::multi_group_by::GroupValuesColumn +/// `GroupValuesRows`: crate::aggregates::group_values::GroupValuesRows +pub fn new_group_values( + schema: SchemaRef, + group_ordering: &GroupOrdering, +) -> Result<Box<dyn GroupValues>> { + if schema.fields.len() == 1 { + let d = schema.fields[0].data_type(); + + macro_rules! downcast_helper { + ($t:ty, $d:ident) => { + return Ok(Box::new(GroupValuesPrimitive::<$t>::new($d.clone()))) + }; + } + + downcast_primitive! { + d => (downcast_helper, d), + _ => {} + } + + match d { + DataType::Date32 => { + downcast_helper!(Date32Type, d); + } + DataType::Date64 => { + downcast_helper!(Date64Type, d); + } + DataType::Time32(t) => match t { + TimeUnit::Second => downcast_helper!(Time32SecondType, d), + TimeUnit::Millisecond => downcast_helper!(Time32MillisecondType, d), + _ => {} + }, + DataType::Time64(t) => match t { + TimeUnit::Microsecond => downcast_helper!(Time64MicrosecondType, d), + TimeUnit::Nanosecond => downcast_helper!(Time64NanosecondType, d), + _ => {} + }, + DataType::Timestamp(t, _tz) => match t { + TimeUnit::Second => downcast_helper!(TimestampSecondType, d), + TimeUnit::Millisecond => downcast_helper!(TimestampMillisecondType, d), + TimeUnit::Microsecond => downcast_helper!(TimestampMicrosecondType, d), + TimeUnit::Nanosecond => downcast_helper!(TimestampNanosecondType, d), + }, + DataType::Decimal128(_, _) => { + downcast_helper!(Decimal128Type, d); + } + DataType::Utf8 => { + return Ok(Box::new(GroupValuesBytes::<i32>::new(OutputType::Utf8))); + } + DataType::LargeUtf8 => { + return Ok(Box::new(GroupValuesBytes::<i64>::new(OutputType::Utf8))); + } + DataType::Utf8View => { + return Ok(Box::new(GroupValuesBytesView::new(OutputType::Utf8View))); + } + DataType::Binary => { + return Ok(Box::new(GroupValuesBytes::<i32>::new(OutputType::Binary))); + } + DataType::LargeBinary => { + return Ok(Box::new(GroupValuesBytes::<i64>::new(OutputType::Binary))); + } + DataType::BinaryView => { + return Ok(Box::new(GroupValuesBytesView::new(OutputType::BinaryView))); + } + DataType::Boolean => { + return Ok(Box::new(GroupValuesBoolean::new())); + } + _ => {} + } + } + + if multi_group_by::supported_schema(schema.as_ref()) { + if matches!(group_ordering, GroupOrdering::None) { + Ok(Box::new(GroupValuesColumn::<false>::try_new(schema)?)) + } else { + Ok(Box::new(GroupValuesColumn::<true>::try_new(schema)?)) + } + } else { + Ok(Box::new(GroupValuesRows::try_new(schema)?)) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs new file mode 100644 index 0000000..5fdbe43 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/boolean.rs
@@ -0,0 +1,493 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use crate::aggregates::group_values::multi_group_by::Nulls; +use crate::aggregates::group_values::multi_group_by::{GroupColumn, nulls_equal_to}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{Array as _, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder}; +use datafusion_common::Result; + +/// An implementation of [`GroupColumn`] for booleans +/// +/// Optimized to skip null buffer construction if the input is known to be non nullable +/// +/// # Template parameters +/// +/// `NULLABLE`: if the data can contain any nulls +#[derive(Debug)] +pub struct BooleanGroupValueBuilder<const NULLABLE: bool> { + buffer: BooleanBufferBuilder, + nulls: MaybeNullBufferBuilder, +} + +impl<const NULLABLE: bool> BooleanGroupValueBuilder<NULLABLE> { + /// Create a new `BooleanGroupValueBuilder` + pub fn new() -> Self { + Self { + buffer: BooleanBufferBuilder::new(0), + nulls: MaybeNullBufferBuilder::new(), + } + } +} + +impl<const NULLABLE: bool> GroupColumn for BooleanGroupValueBuilder<NULLABLE> { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + if NULLABLE { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + } + + self.buffer.get_bit(lhs_row) == array.as_boolean().value(rhs_row) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + if NULLABLE { + if array.is_null(row) { + self.nulls.append(true); + self.buffer.append(bool::default()); + } else { + self.nulls.append(false); + self.buffer.append(array.as_boolean().value(row)); + } + } else { + self.buffer.append(array.as_boolean().value(row)); + } + + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let array = array.as_boolean(); + + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + + if NULLABLE { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + if !result { + equal_to_results.set_bit(idx, false); + } + continue; + } + } + + if self.buffer.get_bit(lhs_row) != array.value(rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let arr = array.as_boolean(); + + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match (NULLABLE, all_null_or_non_null) { + (true, Nulls::Some) => { + for &row in rows { + if array.is_null(row) { + self.nulls.append(true); + self.buffer.append(bool::default()); + } else { + self.nulls.append(false); + self.buffer.append(arr.value(row)); + } + } + } + + (true, Nulls::None) => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.buffer.append(arr.value(row)); + } + } + + (true, Nulls::All) => { + self.nulls.append_n(rows.len(), true); + self.buffer.append_n(rows.len(), bool::default()); + } + + (false, _) => { + for &row in rows { + self.buffer.append(arr.value(row)); + } + } + } + + Ok(()) + } + + fn len(&self) -> usize { + self.buffer.len() + } + + fn size(&self) -> usize { + self.buffer.capacity() / 8 + self.nulls.allocated_size() + } + + fn build(self: Box<Self>) -> ArrayRef { + let Self { mut buffer, nulls } = *self; + + let nulls = nulls.build(); + if !NULLABLE { + assert!(nulls.is_none(), "unexpected nulls in non nullable input"); + } + + let arr = BooleanArray::new(buffer.finish(), nulls); + + Arc::new(arr) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; + + let mut new_builder = BooleanBufferBuilder::new(self.buffer.len()); + new_builder.append_packed_range(n..self.buffer.len(), self.buffer.as_slice()); + std::mem::swap(&mut new_builder, &mut self.buffer); + + // take only first n values from the original builder + new_builder.truncate(n); + + Arc::new(BooleanArray::new(new_builder.finish(), first_n_nulls)) + } +} + +#[cfg(test)] +mod tests { + use arrow::array::{BooleanBufferBuilder, NullBufferBuilder}; + + use super::*; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + #[test] + fn test_nullable_boolean_equal_to() { + let append = |builder: &mut BooleanGroupValueBuilder<true>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &BooleanGroupValueBuilder<true>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_nullable_boolean_equal_to_internal(append, equal_to); + } + + #[test] + fn test_nullable_primitive_vectorized_equal_to() { + let append = |builder: &mut BooleanGroupValueBuilder<true>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &BooleanGroupValueBuilder<true>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_nullable_boolean_equal_to_internal(append, equal_to); + } + + fn test_nullable_boolean_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut BooleanGroupValueBuilder<true>, &ArrayRef, &[usize]), + E: FnMut( + &BooleanGroupValueBuilder<true>, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; values not equal + // - exist not null, input not null; values equal + + // Define BooleanGroupValueBuilder + let mut builder = BooleanGroupValueBuilder::<true>::new(); + let builder_array = Arc::new(BooleanArray::from(vec![ + None, + None, + None, + Some(true), + Some(false), + Some(true), + ])) as ArrayRef; + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5]); + + // Define input array + let (values, _nulls) = BooleanArray::from(vec![ + Some(true), + Some(false), + None, + None, + Some(true), + Some(true), + ]) + .into_parts(); + + // explicitly build a null buffer where one of the null values also happens to match + let mut nulls = NullBufferBuilder::new(6); + nulls.append_non_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_non_null(); + nulls.append_non_null(); + let input_array = Arc::new(BooleanArray::new(values, nulls.finish())) as ArrayRef; + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5], + &input_array, + &[0, 1, 2, 3, 4, 5], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(!results[3]); + assert!(!results[4]); + assert!(results[5]); + } + + #[test] + fn test_not_nullable_primitive_equal_to() { + let append = |builder: &mut BooleanGroupValueBuilder<false>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &BooleanGroupValueBuilder<false>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_not_nullable_boolean_equal_to_internal(append, equal_to); + } + + #[test] + fn test_not_nullable_primitive_vectorized_equal_to() { + let append = |builder: &mut BooleanGroupValueBuilder<false>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &BooleanGroupValueBuilder<false>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_not_nullable_boolean_equal_to_internal(append, equal_to); + } + + fn test_not_nullable_boolean_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut BooleanGroupValueBuilder<false>, &ArrayRef, &[usize]), + E: FnMut( + &BooleanGroupValueBuilder<false>, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - values equal + // - values not equal + + // Define BooleanGroupValueBuilder + let mut builder = BooleanGroupValueBuilder::<false>::new(); + let builder_array = Arc::new(BooleanArray::from(vec![ + Some(false), + Some(true), + Some(false), + Some(true), + ])) as ArrayRef; + append(&mut builder, &builder_array, &[0, 1, 2, 3]); + + // Define input array + let input_array = Arc::new(BooleanArray::from(vec![ + Some(false), + Some(false), + Some(true), + Some(true), + ])) as ArrayRef; + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3], + &input_array, + &[0, 1, 2, 3], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(!results[1]); + assert!(!results[2]); + assert!(results[3]); + } + + #[test] + fn test_nullable_boolean_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = BooleanGroupValueBuilder::<true>::new(); + + // All nulls input array + let all_nulls_input_array = + Arc::new(BooleanArray::from(vec![None, None, None, None, None])) as _; + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = Arc::new(BooleanArray::from(vec![ + Some(false), + Some(true), + Some(false), + Some(true), + Some(true), + ])) as _; + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs new file mode 100644 index 0000000..c83b1da --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs
@@ -0,0 +1,701 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::{ + GroupColumn, Nulls, nulls_equal_to, +}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, BufferBuilder, GenericBinaryArray, + GenericByteArray, GenericStringArray, OffsetSizeTrait, types::GenericStringType, +}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; +use datafusion_common::utils::proxy::VecAllocExt; +use datafusion_common::utils::split_vec_min_alloc; +use datafusion_common::{Result, exec_datafusion_err}; +use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType}; +use std::mem::size_of; +use std::sync::Arc; +use std::vec; + +/// An implementation of [`GroupColumn`] for binary and utf8 types. +/// +/// Stores a collection of binary or utf8 group values in a single buffer +/// in a way that allows: +/// +/// 1. Efficient comparison of incoming rows to existing rows +/// 2. Efficient construction of the final output array +pub struct ByteGroupValueBuilder<O> +where + O: OffsetSizeTrait, +{ + output_type: OutputType, + buffer: BufferBuilder<u8>, + /// Offsets into `buffer` for each distinct value. These offsets as used + /// directly to create the final `GenericBinaryArray`. The `i`th string is + /// stored in the range `offsets[i]..offsets[i+1]` in `buffer`. Null values + /// are stored as a zero length string. + offsets: Vec<O>, + /// Nulls + nulls: MaybeNullBufferBuilder, + /// The maximum size of the buffer for `0` + max_buffer_size: usize, +} + +impl<O> ByteGroupValueBuilder<O> +where + O: OffsetSizeTrait, +{ + pub fn new(output_type: OutputType) -> Self { + Self { + output_type, + buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY), + offsets: vec![O::default()], + nulls: MaybeNullBufferBuilder::new(), + max_buffer_size: if O::IS_LARGE { + i64::MAX as usize + } else { + i32::MAX as usize + }, + } + } + + fn equal_to_inner<B>(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool + where + B: ByteArrayType, + { + let array = array.as_bytes::<B>(); + self.do_equal_to_inner(lhs_row, array, rhs_row) + } + + fn append_val_inner<B>(&mut self, array: &ArrayRef, row: usize) -> Result<()> + where + B: ByteArrayType, + { + let arr = array.as_bytes::<B>(); + if arr.is_null(row) { + self.nulls.append(true); + // nulls need a zero length in the offset buffer + let offset = self.buffer.len(); + self.offsets.push(O::usize_as(offset)); + } else { + self.nulls.append(false); + self.do_append_val_inner(arr, row)?; + } + + Ok(()) + } + + fn vectorized_equal_to_inner<B>( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) where + B: ByteArrayType, + { + let array = array.as_bytes::<B>(); + + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + + if !self.do_equal_to_inner(lhs_row, array, rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append_inner<B>( + &mut self, + array: &ArrayRef, + rows: &[usize], + ) -> Result<()> + where + B: ByteArrayType, + { + let arr = array.as_bytes::<B>(); + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match all_null_or_non_null { + Nulls::Some => { + for &row in rows { + self.append_val_inner::<B>(array, row)? + } + } + + Nulls::None => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.do_append_val_inner(arr, row)?; + } + } + + Nulls::All => { + self.nulls.append_n(rows.len(), true); + + let new_len = self.offsets.len() + rows.len(); + let offset = self.buffer.len(); + self.offsets.resize(new_len, O::usize_as(offset)); + } + } + + Ok(()) + } + + fn do_equal_to_inner<B>( + &self, + lhs_row: usize, + array: &GenericByteArray<B>, + rhs_row: usize, + ) -> bool + where + B: ByteArrayType, + { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + // Otherwise, we need to check their values + self.value(lhs_row) == (array.value(rhs_row).as_ref() as &[u8]) + } + + fn do_append_val_inner<B>( + &mut self, + array: &GenericByteArray<B>, + row: usize, + ) -> Result<()> + where + B: ByteArrayType, + { + let value: &[u8] = array.value(row).as_ref(); + self.buffer.append_slice(value); + + if self.buffer.len() > self.max_buffer_size { + return Err(exec_datafusion_err!( + "offset overflow, buffer size > {}", + self.max_buffer_size + )); + } + + self.offsets.push(O::usize_as(self.buffer.len())); + Ok(()) + } + + /// return the current value of the specified row irrespective of null + pub fn value(&self, row: usize) -> &[u8] { + let l = self.offsets[row].as_usize(); + let r = self.offsets[row + 1].as_usize(); + // Safety: the offsets are constructed correctly and never decrease + unsafe { self.buffer.as_slice().get_unchecked(l..r) } + } +} + +impl<O> GroupColumn for ByteGroupValueBuilder<O> +where + O: OffsetSizeTrait, +{ + fn equal_to(&self, lhs_row: usize, column: &ArrayRef, rhs_row: usize) -> bool { + // Sanity array type + match self.output_type { + OutputType::Binary => { + debug_assert!(matches!( + column.data_type(), + DataType::Binary | DataType::LargeBinary + )); + self.equal_to_inner::<GenericBinaryType<O>>(lhs_row, column, rhs_row) + } + OutputType::Utf8 => { + debug_assert!(matches!( + column.data_type(), + DataType::Utf8 | DataType::LargeUtf8 + )); + self.equal_to_inner::<GenericStringType<O>>(lhs_row, column, rhs_row) + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + } + } + + fn append_val(&mut self, column: &ArrayRef, row: usize) -> Result<()> { + // Sanity array type + match self.output_type { + OutputType::Binary => { + debug_assert!(matches!( + column.data_type(), + DataType::Binary | DataType::LargeBinary + )); + self.append_val_inner::<GenericBinaryType<O>>(column, row)? + } + OutputType::Utf8 => { + debug_assert!(matches!( + column.data_type(), + DataType::Utf8 | DataType::LargeUtf8 + )); + self.append_val_inner::<GenericStringType<O>>(column, row)? + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + }; + + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + // Sanity array type + match self.output_type { + OutputType::Binary => { + debug_assert!(matches!( + array.data_type(), + DataType::Binary | DataType::LargeBinary + )); + self.vectorized_equal_to_inner::<GenericBinaryType<O>>( + lhs_rows, + array, + rhs_rows, + equal_to_results, + ); + } + OutputType::Utf8 => { + debug_assert!(matches!( + array.data_type(), + DataType::Utf8 | DataType::LargeUtf8 + )); + self.vectorized_equal_to_inner::<GenericStringType<O>>( + lhs_rows, + array, + rhs_rows, + equal_to_results, + ); + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + } + } + + fn vectorized_append(&mut self, column: &ArrayRef, rows: &[usize]) -> Result<()> { + match self.output_type { + OutputType::Binary => { + debug_assert!(matches!( + column.data_type(), + DataType::Binary | DataType::LargeBinary + )); + self.vectorized_append_inner::<GenericBinaryType<O>>(column, rows)? + } + OutputType::Utf8 => { + debug_assert!(matches!( + column.data_type(), + DataType::Utf8 | DataType::LargeUtf8 + )); + self.vectorized_append_inner::<GenericStringType<O>>(column, rows)? + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + }; + + Ok(()) + } + + fn len(&self) -> usize { + self.offsets.len() - 1 + } + + fn size(&self) -> usize { + self.buffer.capacity() * size_of::<u8>() + + self.offsets.allocated_size() + + self.nulls.allocated_size() + } + + fn build(self: Box<Self>) -> ArrayRef { + let Self { + output_type, + mut buffer, + offsets, + nulls, + .. + } = *self; + + let null_buffer = nulls.build(); + + // SAFETY: the offsets were constructed correctly in `insert_if_new` -- + // monotonically increasing, overflows were checked. + let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; + let values = buffer.finish(); + match output_type { + OutputType::Binary => { + // SAFETY: the offsets were constructed correctly + Arc::new(unsafe { + GenericBinaryArray::new_unchecked(offsets, values, null_buffer) + }) + } + OutputType::Utf8 => { + // SAFETY: + // 1. the offsets were constructed safely + // + // 2. the input arrays were all the correct type and thus since + // all the values that went in were valid (e.g. utf8) so are all + // the values that come out + Arc::new(unsafe { + GenericStringArray::new_unchecked(offsets, values, null_buffer) + }) + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + } + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(self.len() >= n); + let null_buffer = self.nulls.take_n(n); + let first_remaining_offset = O::as_usize(self.offsets[n]); + + // Given offsets like [0, 2, 4, 5] and n = 1, we expect to get + // offsets [0, 2, 3]. We first create two offsets for first_n as [0, 2] and the remaining as [2, 4, 5]. + // And we shift the offset starting from 0 for the remaining one, [2, 4, 5] -> [0, 2, 3]. + let offset_n = self.offsets[n]; + let mut first_n_offsets = split_vec_min_alloc(&mut self.offsets, n); + // After the split, self.offsets[0] == offset_n in both branches; normalize in-place. + self.offsets.iter_mut().for_each(|o| *o = o.sub(offset_n)); + first_n_offsets.push(offset_n); + + // SAFETY: the offsets were constructed correctly in `insert_if_new` -- + // monotonically increasing, overflows were checked. + let offsets = + unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(first_n_offsets)) }; + + let mut remaining_buffer = + BufferBuilder::new(self.buffer.len() - first_remaining_offset); + // TODO: Current approach copy the remaining and truncate the original one + // Find out a way to avoid copying buffer but split the original one into two. + remaining_buffer.append_slice(&self.buffer.as_slice()[first_remaining_offset..]); + self.buffer.truncate(first_remaining_offset); + let values = self.buffer.finish(); + self.buffer = remaining_buffer; + + match self.output_type { + OutputType::Binary => { + // SAFETY: the offsets were constructed correctly + Arc::new(unsafe { + GenericBinaryArray::new_unchecked(offsets, values, null_buffer) + }) + } + OutputType::Utf8 => { + // SAFETY: + // 1. the offsets were constructed safely + // + // 2. we asserted the input arrays were all the correct type and + // thus since all the values that went in were valid (e.g. utf8) + // so are all the values that come out + Arc::new(unsafe { + GenericStringArray::new_unchecked(offsets, values, null_buffer) + }) + } + _ => unreachable!("View types should use `ArrowBytesViewMap`"), + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::aggregates::group_values::multi_group_by::bytes::ByteGroupValueBuilder; + use arrow::array::{ArrayRef, BooleanBufferBuilder, NullBufferBuilder, StringArray}; + use datafusion_common::DataFusionError; + use datafusion_physical_expr::binary_map::OutputType; + + use super::GroupColumn; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + #[test] + fn test_byte_group_value_builder_overflow() { + let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8); + + let large_string = "a".repeat(1024 * 1024); + + let array = + Arc::new(StringArray::from(vec![Some(large_string.as_str())])) as ArrayRef; + + // Append items until our buffer length is i32::MAX as usize + for _ in 0..2047 { + builder.append_val(&array, 0).unwrap(); + } + + assert!(matches!( + builder.append_val(&array, 0), + Err(DataFusionError::Execution(e)) if e.contains("offset overflow") + )); + + assert_eq!(builder.value(2046), large_string.as_bytes()); + } + + #[test] + fn test_byte_take_n() { + let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8); + let array = Arc::new(StringArray::from(vec![Some("a"), None])) as ArrayRef; + // a, null, null + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 1).unwrap(); + + // (a, null) remaining: null + let output = builder.take_n(2); + assert_eq!(&output, &array); + + // null, a, null, a + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 0).unwrap(); + + // (null, a) remaining: (null, a) + let output = builder.take_n(2); + let array = Arc::new(StringArray::from(vec![None, Some("a")])) as ArrayRef; + assert_eq!(&output, &array); + + let array = Arc::new(StringArray::from(vec![ + Some("a"), + None, + Some("longstringfortest"), + ])) as ArrayRef; + + // null, a, longstringfortest, null, null + builder.append_val(&array, 2).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 1).unwrap(); + + // (null, a, longstringfortest, null) remaining: (null) + let output = builder.take_n(4); + let array = Arc::new(StringArray::from(vec![ + None, + Some("a"), + Some("longstringfortest"), + None, + ])) as ArrayRef; + assert_eq!(&output, &array); + } + + #[test] + fn test_byte_equal_to() { + let append = |builder: &mut ByteGroupValueBuilder<i32>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &ByteGroupValueBuilder<i32>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_byte_equal_to_internal(append, equal_to); + } + + #[test] + fn test_byte_vectorized_equal_to() { + let append = |builder: &mut ByteGroupValueBuilder<i32>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &ByteGroupValueBuilder<i32>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_byte_equal_to_internal(append, equal_to); + } + + #[test] + fn test_byte_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8); + + // All nulls input array + let all_nulls_input_array = Arc::new(StringArray::from(vec![ + Option::<&str>::None, + None, + None, + None, + None, + ])) as _; + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = Arc::new(StringArray::from(vec![ + Some("string1"), + Some("string2"), + Some("string3"), + Some("string4"), + Some("string5"), + ])) as _; + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } + + fn test_byte_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut ByteGroupValueBuilder<i32>, &ArrayRef, &[usize]), + E: FnMut( + &ByteGroupValueBuilder<i32>, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; values not equal + // - exist not null, input not null; values equal + + // Define ByteGroupValueBuilder + let mut builder = ByteGroupValueBuilder::<i32>::new(OutputType::Utf8); + let builder_array = Arc::new(StringArray::from(vec![ + None, + None, + None, + Some("foo"), + Some("bar"), + Some("baz"), + ])) as ArrayRef; + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5]); + + // Define input array + let (offsets, buffer, _nulls) = StringArray::from(vec![ + Some("foo"), + Some("bar"), + None, + None, + Some("foo"), + Some("baz"), + ]) + .into_parts(); + + // explicitly build a null buffer where one of the null values also happens to match + let mut nulls = NullBufferBuilder::new(6); + nulls.append_non_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_non_null(); + nulls.append_non_null(); + let input_array = + Arc::new(StringArray::new(offsets, buffer, nulls.finish())) as ArrayRef; + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5], + &input_array, + &[0, 1, 2, 3, 4, 5], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(!results[3]); + assert!(!results[4]); + assert!(results[5]); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs new file mode 100644 index 0000000..8625772 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs
@@ -0,0 +1,1022 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::{ + GroupColumn, Nulls, nulls_equal_to, +}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, ByteView, GenericByteViewArray, +}; +use arrow::buffer::{Buffer, ScalarBuffer}; +use arrow::datatypes::ByteViewType; +use datafusion_common::Result; +use datafusion_common::utils::split_vec_min_alloc; +use std::marker::PhantomData; +use std::mem::{replace, size_of}; +use std::sync::Arc; + +const BYTE_VIEW_MAX_BLOCK_SIZE: usize = 2 * 1024 * 1024; + +/// An implementation of [`GroupColumn`] for binary view and utf8 view types. +/// +/// Stores a collection of binary view or utf8 view group values in a buffer +/// whose structure is similar to `GenericByteViewArray`, and we can get benefits: +/// +/// 1. Efficient comparison of incoming rows to existing rows +/// 2. Efficient construction of the final output array +/// 3. Efficient to perform `take_n` comparing to use `GenericByteViewBuilder` +pub struct ByteViewGroupValueBuilder<B: ByteViewType> { + /// The views of string values + /// + /// If string len <= 12, the view's format will be: + /// string(12B) | len(4B) + /// + /// If string len > 12, its format will be: + /// offset(4B) | buffer_index(4B) | prefix(4B) | len(4B) + views: Vec<u128>, + + /// The progressing block + /// + /// New values will be inserted into it until its capacity + /// is not enough(detail can see `max_block_size`). + in_progress: Vec<u8>, + + /// The completed blocks + completed: Vec<Buffer>, + + /// The max size of `in_progress` + /// + /// `in_progress` will be flushed into `completed`, and create new `in_progress` + /// when found its remaining capacity(`max_block_size` - `len(in_progress)`), + /// is no enough to store the appended value. + /// + /// Currently it is fixed at 2MB. + max_block_size: usize, + + /// Nulls + nulls: MaybeNullBufferBuilder, + + /// phantom data so the type requires `<B>` + _phantom: PhantomData<B>, +} + +impl<B: ByteViewType> Default for ByteViewGroupValueBuilder<B> { + fn default() -> Self { + Self::new() + } +} + +impl<B: ByteViewType> ByteViewGroupValueBuilder<B> { + pub fn new() -> Self { + Self { + views: Vec::new(), + in_progress: Vec::new(), + completed: Vec::new(), + max_block_size: BYTE_VIEW_MAX_BLOCK_SIZE, + nulls: MaybeNullBufferBuilder::new(), + _phantom: PhantomData {}, + } + } + + /// Set the max block size + fn with_max_block_size(mut self, max_block_size: usize) -> Self { + self.max_block_size = max_block_size; + self + } + + fn equal_to_inner(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + let array = array.as_byte_view::<B>(); + // since this is a single row comparison, don't bother specializing for nulls/buffers + self.do_equal_to_inner::<true, true>(lhs_row, array, rhs_row) + } + + fn append_val_inner(&mut self, array: &ArrayRef, row: usize) { + let arr = array.as_byte_view::<B>(); + + // Null row case, set and return + if arr.is_null(row) { + self.nulls.append(true); + self.views.push(0); + return; + } + + // Not null row case + self.nulls.append(false); + self.do_append_val_inner(arr, row); + } + + // Don't inline to keep the code small and give LLVM the best chance of + // vectorizing the inner loop + #[inline(never)] + fn vectorized_equal_to_inner<const HAS_NULLS: bool, const HAS_BUFFERS: bool>( + &self, + lhs_rows: &[usize], + array: &GenericByteViewArray<B>, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + + if !self.do_equal_to_inner::<HAS_NULLS, HAS_BUFFERS>(lhs_row, array, rhs_row) + { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append_inner( + &mut self, + array: &ArrayRef, + rows: &[usize], + ) -> Result<()> { + let arr = array.as_byte_view::<B>(); + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match all_null_or_non_null { + Nulls::Some => { + for &row in rows { + self.append_val_inner(array, row); + } + } + + Nulls::None => { + self.nulls.append_n(rows.len(), false); + if arr.data_buffers().is_empty() { + // Fast path: all strings are inline (≤12 bytes). + // The input array's u128 views are already in the correct format; + // copy them directly instead of going through value() → make_view(). + self.views.extend(rows.iter().map(|&row| arr.views()[row])); + } else { + // Slow path: some strings may be non-inline (>12 bytes). + // Pre-reserve and delegate to do_append_val_inner which + // reads raw views directly and reuses source prefixes. + self.views.try_reserve(rows.len()).map_err(|e| { + datafusion_common::exec_datafusion_err!( + "failed to reserve {0} views: {e}", + rows.len() + ) + })?; + for &row in rows { + self.do_append_val_inner(arr, row); + } + } + } + + Nulls::All => { + self.nulls.append_n(rows.len(), true); + let new_len = self.views.len() + rows.len(); + self.views.resize(new_len, 0); + } + } + Ok(()) + } + + fn do_append_val_inner(&mut self, array: &GenericByteViewArray<B>, row: usize) + where + B: ByteViewType, + { + // SAFETY: the caller ensures `row` is valid + let view = unsafe { *array.views().get_unchecked(row) }; + let len = view as u32; + + if len <= 12 { + // Inline value: the view is already self-contained, push as-is. + self.views.push(view); + } else { + // Non-inline value: copy the buffer data and construct a new view + // that points into our own buffers, reusing the source prefix. + let src = ByteView::from(view); + self.ensure_in_progress_big_enough(len as usize); + let new_buffer_index = self.completed.len() as u32; + let new_offset = self.in_progress.len() as u32; + let src_buf = &array.data_buffers()[src.buffer_index as usize]; + self.in_progress.extend_from_slice( + &src_buf[src.offset as usize..(src.offset + src.length) as usize], + ); + let new_view = ByteView { + length: src.length, + prefix: src.prefix, + buffer_index: new_buffer_index, + offset: new_offset, + } + .as_u128(); + self.views.push(new_view); + } + } + + fn ensure_in_progress_big_enough(&mut self, value_len: usize) { + debug_assert!(value_len > 12); + let require_cap = self.in_progress.len() + value_len; + + // If current block isn't big enough, flush it and create a new in progress block + if require_cap > self.max_block_size { + let flushed_block = replace( + &mut self.in_progress, + Vec::with_capacity(self.max_block_size), + ); + let buffer = Buffer::from_vec(flushed_block); + self.completed.push(buffer); + } + } + + /// Compare the value at `lhs_row` in this builder with + /// the value at `rhs_row` in input `array` + /// + /// Templated so that the inner compare loop can be + /// specialized based on the input array + #[inline(always)] + fn do_equal_to_inner<const HAS_NULLS: bool, const HAS_BUFFERS: bool>( + &self, + lhs_row: usize, + array: &GenericByteViewArray<B>, + rhs_row: usize, + ) -> bool { + // Check if nulls equal firstly + if HAS_NULLS { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + } + + // Otherwise, we need to check their values + + // SAFETY: the `lhs_row` and rhs_row` are valid + let exist_view = unsafe { *self.views.get_unchecked(lhs_row) }; + let exist_view_len = exist_view as u32; + + let input_view = unsafe { *array.views().get_unchecked(rhs_row) }; + let input_view_len = input_view as u32; + + // fast path, if we know there are no buffers, then the view must be inlined + // so we can simply compare the u128 views + if !HAS_BUFFERS { + return exist_view == input_view; + } + + // The check logic + // - Check len equality + // - If inlined, check inlined value + // - If non-inlined, check prefix and then check value in buffer + // when needed + if exist_view_len != input_view_len { + return false; + } + + if exist_view_len <= 12 { + // both inlined, so compare inlined value + exist_view == input_view + } else { + let exist_prefix = + unsafe { GenericByteViewArray::<B>::inline_value(&exist_view, 4) }; + let input_prefix = + unsafe { GenericByteViewArray::<B>::inline_value(&input_view, 4) }; + + if exist_prefix != input_prefix { + return false; + } + + // get the full values and compare + let exist_full = { + let byte_view = ByteView::from(exist_view); + let buffer_index = byte_view.buffer_index as usize; + let offset = byte_view.offset as usize; + let length = byte_view.length as usize; + debug_assert!(buffer_index <= self.completed.len()); + + unsafe { + if buffer_index < self.completed.len() { + let block = self.completed.get_unchecked(buffer_index); + block.as_slice().get_unchecked(offset..offset + length) + } else { + self.in_progress.get_unchecked(offset..offset + length) + } + } + }; + let input_full: &[u8] = unsafe { array.value_unchecked(rhs_row).as_ref() }; + exist_full == input_full + } + } + + fn build_inner(self) -> ArrayRef { + let Self { + views, + in_progress, + mut completed, + nulls, + .. + } = self; + + // Build nulls + let null_buffer = nulls.build(); + + // Build values + // Flush `in_process` firstly + if !in_progress.is_empty() { + let buffer = Buffer::from(in_progress); + completed.push(buffer); + } + + let views = ScalarBuffer::from(views); + + // Safety: + // * all views were correctly made + // * (if utf8): Input was valid Utf8 so buffer contents are + // valid utf8 as well + unsafe { + Arc::new(GenericByteViewArray::<B>::new_unchecked( + views, + completed, + null_buffer, + )) + } + } + + fn take_n_inner(&mut self, n: usize) -> ArrayRef { + debug_assert!(self.len() >= n); + + // The `n == len` case, we need to take all + if self.len() == n { + let new_builder = Self::new().with_max_block_size(self.max_block_size); + let cur_builder = replace(self, new_builder); + return cur_builder.build_inner(); + } + + // The `n < len` case + // Take n for nulls + let null_buffer = self.nulls.take_n(n); + + // Take n for values: + // - Take first n `view`s from `views` + // + // - Find the last non-inlined `view`, if all inlined, + // we can build array and return happily, otherwise we + // we need to continue to process related buffers + // + // - Get the last related `buffer index`(let's name it `buffer index n`) + // from last non-inlined `view` + // + // - Take buffers, the key is that we need to know if we need to take + // the whole last related buffer. The logic is a bit complex, you can + // detail in `take_buffers_with_whole_last`, `take_buffers_with_partial_last` + // and other related steps in following + // + // - Shift the `buffer index` of remaining non-inlined `views` + // + let first_n_views = split_vec_min_alloc(&mut self.views, n); + + let last_non_inlined_view = first_n_views + .iter() + .rev() + .find(|view| ((**view) as u32) > 12); + + // All taken views inlined + let Some(view) = last_non_inlined_view else { + let views = ScalarBuffer::from(first_n_views); + + // Safety: + // * all views were correctly made + // * (if utf8): Input was valid Utf8 so buffer contents are + // valid utf8 as well + unsafe { + return Arc::new(GenericByteViewArray::<B>::new_unchecked( + views, + Vec::new(), + null_buffer, + )); + } + }; + + // Unfortunately, some taken views non-inlined + let view = ByteView::from(*view); + let last_remaining_buffer_index = view.buffer_index as usize; + + // Check should we take the whole `last_remaining_buffer_index` buffer + let take_whole_last_buffer = self.should_take_whole_buffer( + last_remaining_buffer_index, + (view.offset + view.length) as usize, + ); + + // Take related buffers + let buffers = if take_whole_last_buffer { + self.take_buffers_with_whole_last(last_remaining_buffer_index) + } else { + self.take_buffers_with_partial_last( + last_remaining_buffer_index, + (view.offset + view.length) as usize, + ) + }; + + // Shift `buffer index`s finally + let shifts = if take_whole_last_buffer { + last_remaining_buffer_index + 1 + } else { + last_remaining_buffer_index + }; + + self.views.iter_mut().for_each(|view| { + if (*view as u32) > 12 { + let mut byte_view = ByteView::from(*view); + byte_view.buffer_index -= shifts as u32; + *view = byte_view.as_u128(); + } + }); + + // Build array and return + let views = ScalarBuffer::from(first_n_views); + + // Safety: + // * all views were correctly made + // * (if utf8): Input was valid Utf8 so buffer contents are + // valid utf8 as well + unsafe { + Arc::new(GenericByteViewArray::<B>::new_unchecked( + views, + buffers, + null_buffer, + )) + } + } + + fn take_buffers_with_whole_last( + &mut self, + last_remaining_buffer_index: usize, + ) -> Vec<Buffer> { + if last_remaining_buffer_index == self.completed.len() { + self.flush_in_progress(); + } + self.completed + .drain(0..last_remaining_buffer_index + 1) + .collect() + } + + fn take_buffers_with_partial_last( + &mut self, + last_remaining_buffer_index: usize, + last_take_len: usize, + ) -> Vec<Buffer> { + let mut take_buffers = Vec::with_capacity(last_remaining_buffer_index + 1); + debug_assert!(last_remaining_buffer_index <= self.completed.len()); + + // Process the `last_remaining_buffer_index` buffer before draining so the index is valid. + let last_buffer = if last_remaining_buffer_index < self.completed.len() { + // If it is in `completed`, simply clone + self.completed[last_remaining_buffer_index].clone() + } else { + // If it is `in_progress`, copied `0 ~ offset` part + debug_assert!(last_take_len <= self.in_progress.len()); + let taken_last_buffer = self.in_progress[0..last_take_len].to_vec(); + Buffer::from_vec(taken_last_buffer) + }; + + // Take `0 ~ last_remaining_buffer_index - 1` buffers + if last_remaining_buffer_index > 0 { + take_buffers.extend(self.completed.drain(0..last_remaining_buffer_index)); + } + take_buffers.push(last_buffer); + + take_buffers + } + + #[inline] + fn should_take_whole_buffer(&self, buffer_index: usize, take_len: usize) -> bool { + if buffer_index < self.completed.len() { + take_len == self.completed[buffer_index].len() + } else { + take_len == self.in_progress.len() + } + } + + fn flush_in_progress(&mut self) { + let flushed_block = replace( + &mut self.in_progress, + Vec::with_capacity(self.max_block_size), + ); + let buffer = Buffer::from_vec(flushed_block); + self.completed.push(buffer); + } +} + +impl<B: ByteViewType> GroupColumn for ByteViewGroupValueBuilder<B> { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + self.equal_to_inner(lhs_row, array, rhs_row) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + self.append_val_inner(array, row); + Ok(()) + } + + fn vectorized_equal_to( + &self, + group_indices: &[usize], + array: &ArrayRef, + rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let has_nulls = array.null_count() != 0; + let array = array.as_byte_view::<B>(); + let has_buffers = !array.data_buffers().is_empty(); + // call specialized version based on nulls and buffers presence + match (has_nulls, has_buffers) { + (true, true) => self.vectorized_equal_to_inner::<true, true>( + group_indices, + array, + rows, + equal_to_results, + ), + (true, false) => self.vectorized_equal_to_inner::<true, false>( + group_indices, + array, + rows, + equal_to_results, + ), + (false, true) => self.vectorized_equal_to_inner::<false, true>( + group_indices, + array, + rows, + equal_to_results, + ), + (false, false) => self.vectorized_equal_to_inner::<false, false>( + group_indices, + array, + rows, + equal_to_results, + ), + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + self.vectorized_append_inner(array, rows) + } + + fn len(&self) -> usize { + self.views.len() + } + + fn size(&self) -> usize { + let buffers_size = self + .completed + .iter() + .map(|buf| buf.capacity() * size_of::<u8>()) + .sum::<usize>(); + + self.nulls.allocated_size() + + self.views.capacity() * size_of::<u128>() + + self.in_progress.capacity() * size_of::<u8>() + + buffers_size + + size_of::<Self>() + } + + fn build(self: Box<Self>) -> ArrayRef { + Self::build_inner(*self) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + self.take_n_inner(n) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::aggregates::group_values::multi_group_by::bytes_view::ByteViewGroupValueBuilder; + use arrow::array::{ + ArrayRef, AsArray, BooleanBufferBuilder, NullBufferBuilder, StringViewArray, + }; + use arrow::datatypes::StringViewType; + + use super::GroupColumn; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + #[test] + fn test_byte_view_append_val() { + let mut builder = + ByteViewGroupValueBuilder::<StringViewType>::new().with_max_block_size(60); + let builder_array = StringViewArray::from(vec![ + Some("this string is quite long"), // in buffer 0 + Some("foo"), + None, + Some("bar"), + Some("this string is also quite long"), // buffer 0 + Some("this string is quite long"), // buffer 1 + Some("bar"), + ]); + let builder_array: ArrayRef = Arc::new(builder_array); + for row in 0..builder_array.len() { + builder.append_val(&builder_array, row).unwrap(); + } + + let output = Box::new(builder).build(); + // should be 2 output buffers to hold all the data + assert_eq!(output.as_string_view().data_buffers().len(), 2); + assert_eq!(&output, &builder_array) + } + + #[test] + fn test_byte_view_equal_to() { + let append = |builder: &mut ByteViewGroupValueBuilder<StringViewType>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &ByteViewGroupValueBuilder<StringViewType>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_byte_view_equal_to_internal(append, equal_to); + } + + #[test] + fn test_byte_view_vectorized_equal_to() { + let append = |builder: &mut ByteViewGroupValueBuilder<StringViewType>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &ByteViewGroupValueBuilder<StringViewType>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_byte_view_equal_to_internal(append, equal_to); + } + + #[test] + fn test_byte_view_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = + ByteViewGroupValueBuilder::<StringViewType>::new().with_max_block_size(60); + + // All nulls input array + let all_nulls_input_array = Arc::new(StringViewArray::from(vec![ + Option::<&str>::None, + None, + None, + None, + None, + ])) as _; + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = Arc::new(StringViewArray::from(vec![ + Some("stringview1"), + Some("stringview2"), + Some("stringview3"), + Some("stringview4"), + Some("stringview5"), + ])) as _; + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } + + fn test_byte_view_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut ByteViewGroupValueBuilder<StringViewType>, &ArrayRef, &[usize]), + E: FnMut( + &ByteViewGroupValueBuilder<StringViewType>, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; value lens not equal + // - exist not null, input not null; value not equal(inlined case) + // - exist not null, input not null; value equal(inlined case) + // + // - exist not null, input not null; value not equal + // (non-inlined case + prefix not equal) + // + // - exist not null, input not null; value not equal + // (non-inlined case + value in `completed`) + // + // - exist not null, input not null; value equal + // (non-inlined case + value in `completed`) + // + // - exist not null, input not null; value not equal + // (non-inlined case + value in `in_progress`) + // + // - exist not null, input not null; value equal + // (non-inlined case + value in `in_progress`) + + // Set the block size to 40 for ensuring some unlined values are in `in_progress`, + // and some are in `completed`, so both two branches in `value` function can be covered. + let mut builder = + ByteViewGroupValueBuilder::<StringViewType>::new().with_max_block_size(60); + let builder_array = Arc::new(StringViewArray::from(vec![ + None, + None, + None, + Some("foo"), + Some("bazz"), + Some("foo"), + Some("bar"), + Some("I am a long string for test eq in completed"), + Some("I am a long string for test eq in progress"), + ])) as ArrayRef; + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5, 6, 7, 8]); + + // Define input array + let (views, buffer, _nulls) = StringViewArray::from(vec![ + Some("foo"), + Some("bar"), + None, + None, + Some("baz"), + Some("oof"), + Some("bar"), + Some("i am a long string for test eq in completed"), + Some("I am a long string for test eq in COMPLETED"), + Some("I am a long string for test eq in completed"), + Some("I am a long string for test eq in PROGRESS"), + Some("I am a long string for test eq in progress"), + ]) + .into_parts(); + + // explicitly build a null buffer where one of the null values also happens to match + let mut nulls = NullBufferBuilder::new(9); + nulls.append_non_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_non_null(); + let input_array = + Arc::new(StringViewArray::new(views, buffer, nulls.finish())) as ArrayRef; + + // Check + let mut equal_to_results = make_true_buffer(input_array.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5, 6, 7, 7, 7, 8, 8], + &input_array, + &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(!results[3]); + assert!(!results[4]); + assert!(!results[5]); + assert!(results[6]); + assert!(!results[7]); + assert!(!results[8]); + assert!(results[9]); + assert!(!results[10]); + assert!(results[11]); + } + + #[test] + fn test_byte_view_take_n() { + // ####### Define cases and init ####### + + // `take_n` is really complex, we should consider and test following situations: + // 1. Take nulls + // 2. Take all `inlined`s + // 3. Take non-inlined + partial last buffer in `completed` + // 4. Take non-inlined + whole last buffer in `completed` + // 5. Take non-inlined + partial last `in_progress` + // 6. Take non-inlined + whole last buffer in `in_progress` + // 7. Take all views at once + + let mut builder = + ByteViewGroupValueBuilder::<StringViewType>::new().with_max_block_size(60); + let input_array = StringViewArray::from(vec![ + // Test situation 1 + None, + None, + // Test situation 2 (also test take null together) + None, + Some("foo"), + Some("bar"), + // Test situation 3 (also test take null + inlined) + None, + Some("foo"), + Some("this string is quite long"), + Some("this string is also quite long"), + // Test situation 4 (also test take null + inlined) + None, + Some("bar"), + Some("this string is quite long"), + // Test situation 5 (also test take null + inlined) + None, + Some("foo"), + Some("another string that is is quite long"), + Some("this string not so long"), + // Test situation 6 (also test take null + inlined + insert again after taking) + None, + Some("bar"), + Some("this string is quite long"), + // Insert 4 and just take 3 to ensure it will go the path of situation 6 + None, + // Finally, we create a new builder, insert the whole array and then + // take whole at once for testing situation 7 + ]); + + let input_array: ArrayRef = Arc::new(input_array); + let first_ones_to_append = 16; // For testing situation 1~5 + let second_ones_to_append = 4; // For testing situation 6 + let final_ones_to_append = input_array.len(); // For testing situation 7 + + // ####### Test situation 1~5 ####### + for row in 0..first_ones_to_append { + builder.append_val(&input_array, row).unwrap(); + } + + assert_eq!(builder.completed.len(), 2); + assert_eq!(builder.in_progress.len(), 59); + + // Situation 1 + let taken_array = builder.take_n(2); + assert_eq!(&taken_array, &input_array.slice(0, 2)); + + // Situation 2 + let taken_array = builder.take_n(3); + assert_eq!(&taken_array, &input_array.slice(2, 3)); + + // Situation 3 + let taken_array = builder.take_n(3); + assert_eq!(&taken_array, &input_array.slice(5, 3)); + + let taken_array = builder.take_n(1); + assert_eq!(&taken_array, &input_array.slice(8, 1)); + + // Situation 4 + let taken_array = builder.take_n(3); + assert_eq!(&taken_array, &input_array.slice(9, 3)); + + // Situation 5 + let taken_array = builder.take_n(3); + assert_eq!(&taken_array, &input_array.slice(12, 3)); + + let taken_array = builder.take_n(1); + assert_eq!(&taken_array, &input_array.slice(15, 1)); + + // ####### Test situation 6 ####### + assert!(builder.completed.is_empty()); + assert!(builder.in_progress.is_empty()); + assert!(builder.views.is_empty()); + + for row in first_ones_to_append..first_ones_to_append + second_ones_to_append { + builder.append_val(&input_array, row).unwrap(); + } + + assert!(builder.completed.is_empty()); + assert_eq!(builder.in_progress.len(), 25); + + let taken_array = builder.take_n(3); + assert_eq!(&taken_array, &input_array.slice(16, 3)); + + // ####### Test situation 7 ####### + // Create a new builder + let mut builder = + ByteViewGroupValueBuilder::<StringViewType>::new().with_max_block_size(60); + + for row in 0..final_ones_to_append { + builder.append_val(&input_array, row).unwrap(); + } + + assert_eq!(builder.completed.len(), 3); + assert_eq!(builder.in_progress.len(), 25); + + let taken_array = builder.take_n(final_ones_to_append); + assert_eq!(&taken_array, &input_array); + } + + #[test] + fn test_byte_view_take_n_partial_completed_nonzero_index() { + let mut builder = + ByteViewGroupValueBuilder::<StringViewType>::new().with_max_block_size(30); + let input_array = StringViewArray::from(vec![ + Some("aaaaaaaaaaaaaa"), + Some("bbbbbbbbbbbbbb"), + Some("cccccccccccccc"), + Some("dddddddddddddd"), + Some("eeeeeeeeeeeeee"), + ]); + let input_array: ArrayRef = Arc::new(input_array); + + for row in 0..input_array.len() { + builder.append_val(&input_array, row).unwrap(); + } + + assert_eq!(builder.completed.len(), 2); + assert_eq!(builder.in_progress.len(), 14); + + let taken_array = builder.take_n(3); + assert_eq!(&taken_array, &input_array.slice(0, 3)); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs new file mode 100644 index 0000000..589083c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs
@@ -0,0 +1,515 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::{ + GroupColumn, Nulls, nulls_equal_to, +}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeBinaryArray, +}; +use arrow::buffer::{Buffer, NullBuffer}; +use datafusion_common::utils::proxy::VecAllocExt; +use datafusion_common::utils::split_vec_min_alloc; +use datafusion_common::{Result, exec_datafusion_err}; +use std::sync::Arc; + +/// An implementation of [`GroupColumn`] for `FixedSizeBinary` values +/// +/// Stores the group values in a single flat buffer, `byte_width` bytes per +/// value, in a way that allows: +/// +/// 1. Efficient comparison of incoming rows to existing rows +/// 2. Efficient construction of the final output array (the buffer is handed +/// to [`FixedSizeBinaryArray`] as-is, no offsets needed) +/// +/// Null values occupy `byte_width` zeroed bytes in the buffer so that the +/// value of row `i` is always stored at `i * byte_width..(i + 1) * byte_width`. +pub struct FixedSizeBinaryGroupValueBuilder { + /// The width in bytes of each value, from `DataType::FixedSizeBinary` + byte_width: usize, + /// The flattened group values, `byte_width` bytes per value + buffer: Vec<u8>, + /// The number of group values stored + /// + /// Tracked explicitly rather than derived from `buffer.len()` because + /// `byte_width` may be `0` + len: usize, + /// Null state (null rows still occupy `byte_width` bytes in `buffer`) + nulls: MaybeNullBufferBuilder, +} + +impl FixedSizeBinaryGroupValueBuilder { + /// Create a new builder for values of `byte_width` bytes each + /// + /// `byte_width` is the width carried by `DataType::FixedSizeBinary` and + /// must be non-negative (negative widths are rejected by the dispatch in + /// `make_group_column`) + pub fn new(byte_width: i32) -> Self { + debug_assert!(byte_width >= 0); + Self { + byte_width: byte_width as usize, + buffer: Vec::new(), + len: 0, + nulls: MaybeNullBufferBuilder::new(), + } + } + + fn do_append_val_inner(&mut self, array: &FixedSizeBinaryArray, row: usize) { + if array.is_null(row) { + self.nulls.append(true); + // Null rows still occupy `byte_width` (zeroed) bytes in the + // buffer so the value offset stays a function of the row index + self.buffer.resize(self.buffer.len() + self.byte_width, 0); + } else { + self.nulls.append(false); + self.buffer.extend_from_slice(array.value(row)); + } + self.len += 1; + } + + fn do_equal_to_inner( + &self, + lhs_row: usize, + array: &FixedSizeBinaryArray, + rhs_row: usize, + ) -> bool { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + // Otherwise, we need to check their values + self.value(lhs_row) == array.value(rhs_row) + } + + /// return the current value of the specified row irrespective of null + /// (null rows store `byte_width` zeroed bytes) + pub fn value(&self, row: usize) -> &[u8] { + let start = row * self.byte_width; + &self.buffer[start..start + self.byte_width] + } + + /// Assemble an output array from `values` + `nulls` parts + /// + /// Uses `try_new_with_len` rather than `try_new` because the length + /// cannot be derived from the values buffer when `byte_width == 0` + fn build_array( + byte_width: usize, + values: Vec<u8>, + nulls: Option<NullBuffer>, + len: usize, + ) -> ArrayRef { + let array = FixedSizeBinaryArray::try_new_with_len( + byte_width as i32, + Buffer::from(values), + nulls, + len, + ) + .expect("buffer, nulls and len kept consistent on append"); + Arc::new(array) + } +} + +impl GroupColumn for FixedSizeBinaryGroupValueBuilder { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + self.do_equal_to_inner(lhs_row, array.as_fixed_size_binary(), rhs_row) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let arr = array.as_fixed_size_binary(); + debug_assert_eq!(arr.value_size(), self.byte_width); + self.do_append_val_inner(arr, row); + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let array = array.as_fixed_size_binary(); + + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + // Has found not equal to in previous column, don't need to check + if !equal_to_results.get_bit(idx) { + continue; + } + + if !self.do_equal_to_inner(lhs_row, array, rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let arr = array.as_fixed_size_binary(); + debug_assert_eq!(arr.value_size(), self.byte_width); + + let reserve_bytes = rows.len() * self.byte_width; + self.buffer.try_reserve(reserve_bytes).map_err(|e| { + exec_datafusion_err!("failed to reserve {reserve_bytes} bytes: {e}") + })?; + + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match all_null_or_non_null { + Nulls::Some => { + for &row in rows { + self.do_append_val_inner(arr, row); + } + } + + Nulls::None => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.buffer.extend_from_slice(arr.value(row)); + } + self.len += rows.len(); + } + + Nulls::All => { + self.nulls.append_n(rows.len(), true); + self.buffer + .resize(self.buffer.len() + rows.len() * self.byte_width, 0); + self.len += rows.len(); + } + } + + Ok(()) + } + + fn len(&self) -> usize { + self.len + } + + fn size(&self) -> usize { + self.buffer.allocated_size() + self.nulls.allocated_size() + } + + fn build(self: Box<Self>) -> ArrayRef { + let Self { + byte_width, + buffer, + len, + nulls, + } = *self; + + Self::build_array(byte_width, buffer, nulls.build(), len) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(self.len >= n); + + let null_buffer = self.nulls.take_n(n); + let first_n = split_vec_min_alloc(&mut self.buffer, n * self.byte_width); + self.len -= n; + + Self::build_array(self.byte_width, first_n, null_buffer, n) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::aggregates::group_values::multi_group_by::fixed_size_binary::FixedSizeBinaryGroupValueBuilder; + use arrow::array::{ArrayRef, BooleanBufferBuilder, FixedSizeBinaryArray}; + + use super::GroupColumn; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + fn make_array(values: Vec<Option<&[u8]>>, byte_width: i32) -> ArrayRef { + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.into_iter(), + byte_width, + ) + .unwrap(), + ) + } + + #[test] + fn test_fixed_size_binary_equal_to() { + let append = |builder: &mut FixedSizeBinaryGroupValueBuilder, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &FixedSizeBinaryGroupValueBuilder, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_fixed_size_binary_equal_to_internal(append, equal_to); + } + + #[test] + fn test_fixed_size_binary_vectorized_equal_to() { + let append = |builder: &mut FixedSizeBinaryGroupValueBuilder, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &FixedSizeBinaryGroupValueBuilder, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_fixed_size_binary_equal_to_internal(append, equal_to); + } + + fn test_fixed_size_binary_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut FixedSizeBinaryGroupValueBuilder, &ArrayRef, &[usize]), + E: FnMut( + &FixedSizeBinaryGroupValueBuilder, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; values not equal + // - exist not null, input not null; values equal + + // Define FixedSizeBinaryGroupValueBuilder + let mut builder = FixedSizeBinaryGroupValueBuilder::new(3); + let builder_array = make_array( + vec![ + None, + None, + None, + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + Some(b"baz".as_slice()), + ], + 3, + ); + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5]); + + // Define input array; the value behind the null at row 3 happens to + // match the existing group value to make sure nulls win over values + let input_array = make_array( + vec![ + Some(b"foo".as_slice()), + None, + None, + None, + Some(b"foo".as_slice()), + Some(b"baz".as_slice()), + ], + 3, + ); + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5], + &input_array, + &[0, 1, 2, 3, 4, 5], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(!results[3]); + assert!(!results[4]); + assert!(results[5]); + } + + #[test] + fn test_fixed_size_binary_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + + // All nulls input array + let all_nulls_input_array = make_array(vec![None, None, None, None, None], 2); + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = make_array( + vec![ + Some(b"v1".as_slice()), + Some(b"v2".as_slice()), + Some(b"v3".as_slice()), + Some(b"v4".as_slice()), + Some(b"v5".as_slice()), + ], + 2, + ); + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } + + #[test] + fn test_fixed_size_binary_take_n() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let array = make_array(vec![Some(b"aa".as_slice()), None], 2); + // aa, null, null + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 1).unwrap(); + + // (aa, null) remaining: null + let output = builder.take_n(2); + assert_eq!(&output, &array); + assert_eq!(builder.len(), 1); + + // null, aa, null, aa + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 0).unwrap(); + + // (null, aa) remaining: (null, aa) + let output = builder.take_n(2); + let expected = make_array(vec![None, Some(b"aa".as_slice())], 2); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 2); + + // take the remaining (null, aa) + let output = builder.take_n(2); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 0); + } + + #[test] + fn test_fixed_size_binary_build() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let array = make_array( + vec![Some(b"aa".as_slice()), None, Some(b"bb".as_slice())], + 2, + ); + builder.vectorized_append(&array, &[0, 1, 2]).unwrap(); + assert_eq!(builder.len(), 3); + + let output = Box::new(builder).build(); + assert_eq!(&output, &array); + } + + #[test] + fn test_zero_width_fixed_size_binary() { + // A zero byte width is valid per the Arrow spec; the builder must + // track its length without relying on the (empty) values buffer + let mut builder = FixedSizeBinaryGroupValueBuilder::new(0); + let array = make_array(vec![Some(b"".as_slice()), None, Some(b"".as_slice())], 0); + + builder.vectorized_append(&array, &[0, 1, 2]).unwrap(); + assert_eq!(builder.len(), 3); + + // Empty values compare equal, null only equals null + assert!(builder.equal_to(0, &array, 2)); + assert!(builder.equal_to(1, &array, 1)); + assert!(!builder.equal_to(1, &array, 0)); + + let output = builder.take_n(2); + let expected = make_array(vec![Some(b"".as_slice()), None], 0); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 1); + + let output = Box::new(builder).build(); + let expected = make_array(vec![Some(b"".as_slice())], 0); + assert_eq!(&output, &expected); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs new file mode 100644 index 0000000..5b474f3 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
@@ -0,0 +1,2536 @@ +// 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. + +//! `GroupValues` implementations for multi group by cases + +mod boolean; +mod bytes; +pub mod bytes_view; +mod fixed_size_binary; +pub mod primitive; +pub mod row_backed; + +use std::mem::{self, size_of}; + +use crate::aggregates::group_values::GroupValues; +use crate::aggregates::group_values::multi_group_by::{ + boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder, + bytes_view::ByteViewGroupValueBuilder, + fixed_size_binary::FixedSizeBinaryGroupValueBuilder, + primitive::PrimitiveGroupValueBuilder, row_backed::RowsGroupColumn, +}; +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::compute::cast; +use arrow::datatypes::{ + BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, + DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, + DurationSecondType, Field, Float16Type, Float32Type, Float64Type, Int8Type, + Int16Type, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, + IntervalUnit, IntervalYearMonthType, Schema, SchemaRef, StringViewType, + Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, + TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, + UInt64Type, +}; +use datafusion_common::hash_utils::RandomState; +use datafusion_common::hash_utils::create_hashes; +use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_expr::EmitTo; +use datafusion_physical_expr::binary_map::OutputType; + +use hashbrown::hash_table::HashTable; + +const NON_INLINED_FLAG: u64 = 0x8000000000000000; +const VALUE_MASK: u64 = 0x7FFFFFFFFFFFFFFF; + +/// Trait for storing a single column of group values in [`GroupValuesColumn`] +/// +/// Implementations of this trait store an in-progress collection of group values +/// (similar to various builders in Arrow-rs) that allow for quick comparison to +/// incoming rows. +/// +/// [`GroupValuesColumn`]: crate::aggregates::group_values::GroupValuesColumn +pub trait GroupColumn: Send + Sync { + /// Returns equal if the row stored in this builder at `lhs_row` is equal to + /// the row in `array` at `rhs_row` + /// + /// Note that this comparison returns true if both elements are NULL + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool; + + /// Appends the row at `row` in `array` to this builder + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()>; + + /// The vectorized version equal to + /// + /// When found nth row stored in this builder at `lhs_row` + /// is equal to the row in `array` at `rhs_row`, + /// it will record the `true` result at the corresponding + /// position in `equal_to_results`. + /// + /// And if found nth result in `equal_to_results` is already + /// `false`, the check for nth row will be skipped. + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ); + + /// The vectorized version `append_val` + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()>; + + /// Returns the number of rows stored in this builder + fn len(&self) -> usize; + + /// true if len == 0 + fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Returns the number of bytes used by this [`GroupColumn`] + fn size(&self) -> usize; + + /// Builds a new array from all of the stored rows + fn build(self: Box<Self>) -> ArrayRef; + + /// Builds a new array from the first `n` stored rows, shifting the + /// remaining rows to the start of the builder + fn take_n(&mut self, n: usize) -> ArrayRef; +} + +/// Determines if the nullability of the existing and new input array can be used +/// to short-circuit the comparison of the two values. +/// +/// Returns `Some(result)` if the result of the comparison can be determined +/// from the nullness of the two values, and `None` if the comparison must be +/// done on the values themselves. +pub fn nulls_equal_to(lhs_null: bool, rhs_null: bool) -> Option<bool> { + match (lhs_null, rhs_null) { + (true, true) => Some(true), + (false, true) | (true, false) => Some(false), + _ => None, + } +} + +/// The view of indices pointing to the actual values in `GroupValues` +/// +/// If only single `group index` represented by view, +/// value of view is just the `group index`, and we call it a `inlined view`. +/// +/// If multiple `group indices` represented by view, +/// value of view is the actually the index pointing to `group indices`, +/// and we call it `non-inlined view`. +/// +/// The view(a u64) format is like: +/// +---------------------+---------------------------------------------+ +/// | inlined flag(1bit) | group index / index to group indices(63bit) | +/// +---------------------+---------------------------------------------+ +/// +/// `inlined flag`: 1 represents `non-inlined`, and 0 represents `inlined` +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct GroupIndexView(u64); + +impl GroupIndexView { + #[inline] + pub fn is_non_inlined(&self) -> bool { + (self.0 & NON_INLINED_FLAG) > 0 + } + + #[inline] + pub fn new_inlined(group_index: u64) -> Self { + Self(group_index) + } + + #[inline] + pub fn new_non_inlined(list_offset: u64) -> Self { + let non_inlined_value = list_offset | NON_INLINED_FLAG; + Self(non_inlined_value) + } + + #[inline] + pub fn value(&self) -> u64 { + self.0 & VALUE_MASK + } +} + +/// A [`GroupValues`] that stores multiple columns of group values, +/// and supports vectorized operators for them +pub struct GroupValuesColumn<const STREAMING: bool> { + /// The output schema + schema: SchemaRef, + + /// Logically maps group values to a group_index in + /// [`Self::group_values`] and in each accumulator + /// + /// It is a `hashtable` based on `hashbrown`. + /// + /// Key and value in the `hashtable`: + /// - The `key` is `hash value(u64)` of the `group value` + /// - The `value` is the `group values` with the same `hash value` + /// + /// We don't really store the actual `group values` in `hashtable`, + /// instead we store the `group indices` pointing to values in `GroupValues`. + /// And we use [`GroupIndexView`] to represent such `group indices` in table. + /// + map: HashTable<(u64, GroupIndexView)>, + + /// The size of `map` in bytes + map_size: usize, + + /// The lists for group indices with the same hash value + /// + /// It is possible that hash value collision exists, + /// and we will chain the `group indices` with same hash value + /// + /// The chained indices is like: + /// `latest group index -> older group index -> even older group index -> ...` + group_index_lists: Vec<Vec<usize>>, + + /// When emitting first n, we need to decrease/erase group indices in + /// `map` and `group_index_lists`. + /// + /// This buffer is used to temporarily store the remaining group indices in + /// a specific list in `group_index_lists`. + emit_group_index_list_buffer: Vec<usize>, + + /// Buffers for `vectorized_append` and `vectorized_equal_to` + vectorized_operation_buffers: VectorizedOperationBuffers, + + /// The actual group by values, stored column-wise. Compare from + /// the left to right, each column is stored as [`GroupColumn`]. + /// + /// Performance tests showed that this design is faster than using the + /// more general purpose [`GroupValuesRows`]. See the ticket for details: + /// <https://github.com/apache/datafusion/pull/12269> + /// + /// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows + group_values: Vec<Box<dyn GroupColumn>>, + + /// reused buffer to store hashes + hashes_buffer: Vec<u64>, + + /// Random state for creating hashes + random_state: RandomState, +} + +/// Buffers to store intermediate results in `vectorized_append` +/// and `vectorized_equal_to`, for reducing memory allocation +struct VectorizedOperationBuffers { + /// The `vectorized append` row indices buffer + append_row_indices: Vec<usize>, + + /// The `vectorized_equal_to` row indices buffer + equal_to_row_indices: Vec<usize>, + + /// The `vectorized_equal_to` group indices buffer + equal_to_group_indices: Vec<usize>, + + /// The `vectorized_equal_to` result buffer (bitmask) + equal_to_results: BooleanBufferBuilder, + + /// The buffer for storing row indices found not equal to + /// exist groups in `group_values` in `vectorized_equal_to`. + /// We will perform `scalarized_intern` for such rows. + remaining_row_indices: Vec<usize>, +} + +impl Default for VectorizedOperationBuffers { + fn default() -> Self { + Self { + append_row_indices: Vec::new(), + equal_to_row_indices: Vec::new(), + equal_to_group_indices: Vec::new(), + equal_to_results: BooleanBufferBuilder::new(0), + remaining_row_indices: Vec::new(), + } + } +} + +impl VectorizedOperationBuffers { + fn clear(&mut self) { + self.append_row_indices.clear(); + self.equal_to_row_indices.clear(); + self.equal_to_group_indices.clear(); + self.remaining_row_indices.clear(); + } +} + +impl<const STREAMING: bool> GroupValuesColumn<STREAMING> { + // ======================================================================== + // Initialization functions + // ======================================================================== + + /// Create a new instance of GroupValuesColumn if supported for the specified schema + pub fn try_new(schema: SchemaRef) -> Result<Self> { + let map = HashTable::with_capacity(0); + let group_values = Self::build_group_columns(&schema)?; + Ok(Self { + schema, + map, + group_index_lists: Vec::new(), + emit_group_index_list_buffer: Vec::new(), + vectorized_operation_buffers: VectorizedOperationBuffers::default(), + map_size: 0, + group_values, + hashes_buffer: Default::default(), + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + }) + } + + /// Build one fresh [`GroupColumn`] per field in the schema. + /// + /// Used at construction time (`try_new`) and to repopulate the column + /// vector after operations that drain it (`emit(EmitTo::All)`, + /// `clear_shrink`). Centralising it keeps the post-condition that + /// `self.group_values` always contains exactly one builder per schema + /// field outside of those transient drain points. + fn build_group_columns(schema: &Schema) -> Result<Vec<Box<dyn GroupColumn>>> { + let mut v: Vec<Box<dyn GroupColumn>> = Vec::with_capacity(schema.fields().len()); + for f in schema.fields().iter() { + v.push(make_group_column(f.as_ref())?); + } + Ok(v) + } + + // ======================================================================== + // Scalarized intern + // ======================================================================== + + /// Scalarized intern + /// + /// This is used only for `streaming aggregation`, because `streaming aggregation` + /// depends on the order between `input rows` and their corresponding `group indices`. + /// + /// For example, assuming `input rows` in `cols` with 4 new rows + /// (not equal to `exist rows` in `group_values`, and need to create + /// new groups for them): + /// + /// ```text + /// row1 (hash collision with the exist rows) + /// row2 + /// row3 (hash collision with the exist rows) + /// row4 + /// ``` + /// + /// # In `scalarized_intern`, their `group indices` will be + /// + /// ```text + /// row1 --> 0 + /// row2 --> 1 + /// row3 --> 2 + /// row4 --> 3 + /// ``` + /// + /// `Group indices` order agrees with their input order, and the `streaming aggregation` + /// depends on this. + /// + /// # However In `vectorized_intern`, their `group indices` will be + /// + /// ```text + /// row1 --> 2 + /// row2 --> 0 + /// row3 --> 3 + /// row4 --> 1 + /// ``` + /// + /// `Group indices` order are against with their input order, and this will lead to error + /// in `streaming aggregation`. + fn scalarized_intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec<usize>, + ) -> Result<()> { + let n_rows = cols[0].len(); + + // tracks to which group each of the input rows belongs + groups.clear(); + + // 1.1 Calculate the group keys for the group values + let batch_hashes = &mut self.hashes_buffer; + batch_hashes.clear(); + batch_hashes.resize(n_rows, 0); + create_hashes(cols, &self.random_state, batch_hashes)?; + + for (row, &target_hash) in batch_hashes.iter().enumerate() { + let entry = self + .map + .find_mut(target_hash, |(exist_hash, group_idx_view)| { + // It is ensured to be inlined in `scalarized_intern` + debug_assert!(!group_idx_view.is_non_inlined()); + + // Somewhat surprisingly, this closure can be called even if the + // hash doesn't match, so check the hash first with an integer + // comparison first avoid the more expensive comparison with + // group value. https://github.com/apache/datafusion/pull/11718 + if target_hash != *exist_hash { + return false; + } + + fn check_row_equal( + array_row: &dyn GroupColumn, + lhs_row: usize, + array: &ArrayRef, + rhs_row: usize, + ) -> bool { + array_row.equal_to(lhs_row, array, rhs_row) + } + + for (i, group_val) in self.group_values.iter().enumerate() { + if !check_row_equal( + group_val.as_ref(), + group_idx_view.value() as usize, + &cols[i], + row, + ) { + return false; + } + } + + true + }); + + let group_idx = match entry { + // Existing group_index for this group value + Some((_hash, group_idx_view)) => group_idx_view.value() as usize, + // 1.2 Need to create new entry for the group + None => { + // Add new entry to aggr_state and save newly created index + // let group_idx = group_values.num_rows(); + // group_values.push(group_rows.row(row)); + + let mut checklen = 0; + let group_idx = self.group_values[0].len(); + for (i, group_value) in self.group_values.iter_mut().enumerate() { + group_value.append_val(&cols[i], row)?; + let len = group_value.len(); + if i == 0 { + checklen = len; + } else { + debug_assert_eq!(checklen, len); + } + } + + // for hasher function, use precomputed hash value + self.map.insert_accounted( + (target_hash, GroupIndexView::new_inlined(group_idx as u64)), + |(hash, _group_index)| *hash, + &mut self.map_size, + ); + group_idx + } + }; + groups.push(group_idx); + } + + Ok(()) + } + + // ======================================================================== + // Vectorized intern + // ======================================================================== + + /// Vectorized intern + /// + /// This is used in `non-streaming aggregation` without requiring the order between + /// rows in `cols` and corresponding groups in `group_values`. + /// + /// The vectorized approach can offer higher performance for avoiding row by row + /// downcast for `cols` and being able to implement even more optimizations(like simd). + fn vectorized_intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec<usize>, + ) -> Result<()> { + let n_rows = cols[0].len(); + + // tracks to which group each of the input rows belongs + groups.clear(); + groups.resize(n_rows, usize::MAX); + + let mut batch_hashes = mem::take(&mut self.hashes_buffer); + batch_hashes.clear(); + batch_hashes.resize(n_rows, 0); + create_hashes(cols, &self.random_state, &mut batch_hashes)?; + + // General steps for one round `vectorized equal_to & append`: + // 1. Collect vectorized context by checking hash values of `cols` in `map`, + // mainly fill `vectorized_append_row_indices`, `vectorized_equal_to_row_indices` + // and `vectorized_equal_to_group_indices` + // + // 2. Perform `vectorized_append` for `vectorized_append_row_indices`. + // `vectorized_append` must be performed before `vectorized_equal_to`, + // because some `group indices` in `vectorized_equal_to_group_indices` + // maybe still point to no actual values in `group_values` before performing append. + // + // 3. Perform `vectorized_equal_to` for `vectorized_equal_to_row_indices` + // and `vectorized_equal_to_group_indices`. If found some rows in input `cols` + // not equal to `exist rows` in `group_values`, place them in `remaining_row_indices` + // and perform `scalarized_intern_remaining` for them similar as `scalarized_intern` + // after. + // + // 4. Perform `scalarized_intern_remaining` for rows mentioned above, about in what situation + // we will process this can see the comments of `scalarized_intern_remaining`. + // + + // 1. Collect vectorized context by checking hash values of `cols` in `map` + self.collect_vectorized_process_context(&batch_hashes, groups); + + // 2. Perform `vectorized_append` + self.vectorized_append(cols)?; + + // 3. Perform `vectorized_equal_to` + self.vectorized_equal_to(cols, groups); + + // 4. Perform scalarized inter for remaining rows + // (about remaining rows, can see comments for `remaining_row_indices`) + self.scalarized_intern_remaining(cols, &batch_hashes, groups)?; + + self.hashes_buffer = batch_hashes; + + Ok(()) + } + + /// Collect vectorized context by checking hash values of `cols` in `map` + /// + /// 1. If bucket not found + /// - Build and insert the `new inlined group index view` + /// and its hash value to `map` + /// - Add row index to `vectorized_append_row_indices` + /// - Set group index to row in `groups` + /// + /// 2. bucket found + /// - Add row index to `vectorized_equal_to_row_indices` + /// - Check if the `group index view` is `inlined` or `non_inlined`: + /// If it is inlined, add to `vectorized_equal_to_group_indices` directly. + /// Otherwise get all group indices from `group_index_lists`, and add them. + fn collect_vectorized_process_context( + &mut self, + batch_hashes: &[u64], + groups: &mut [usize], + ) { + self.vectorized_operation_buffers.append_row_indices.clear(); + self.vectorized_operation_buffers + .equal_to_row_indices + .clear(); + self.vectorized_operation_buffers + .equal_to_group_indices + .clear(); + + for (row, &target_hash) in batch_hashes.iter().enumerate() { + let entry = self + .map + .find(target_hash, |(exist_hash, _)| target_hash == *exist_hash); + + let Some((_, group_index_view)) = entry else { + // 1. Bucket not found case + // Build `new inlined group index view` + let current_group_idx = self.group_values[0].len() + + self.vectorized_operation_buffers.append_row_indices.len(); + let group_index_view = + GroupIndexView::new_inlined(current_group_idx as u64); + + // Insert the `group index view` and its hash into `map` + // for hasher function, use precomputed hash value + self.map.insert_accounted( + (target_hash, group_index_view), + |(hash, _)| *hash, + &mut self.map_size, + ); + + // Add row index to `vectorized_append_row_indices` + self.vectorized_operation_buffers + .append_row_indices + .push(row); + + // Set group index to row in `groups` + groups[row] = current_group_idx; + + continue; + }; + + // 2. bucket found + // Check if the `group index view` is `inlined` or `non_inlined` + if group_index_view.is_non_inlined() { + // Non-inlined case, the value of view is offset in `group_index_lists`. + // We use it to get `group_index_list`, and add related `rows` and `group_indices` + // into `vectorized_equal_to_row_indices` and `vectorized_equal_to_group_indices`. + let list_offset = group_index_view.value() as usize; + let group_index_list = &self.group_index_lists[list_offset]; + + self.vectorized_operation_buffers + .equal_to_group_indices + .extend_from_slice(group_index_list); + self.vectorized_operation_buffers + .equal_to_row_indices + .extend(std::iter::repeat_n(row, group_index_list.len())); + } else { + let group_index = group_index_view.value() as usize; + self.vectorized_operation_buffers + .equal_to_row_indices + .push(row); + self.vectorized_operation_buffers + .equal_to_group_indices + .push(group_index); + } + } + } + + /// Perform `vectorized_append`` for `rows` in `vectorized_append_row_indices` + fn vectorized_append(&mut self, cols: &[ArrayRef]) -> Result<()> { + if self + .vectorized_operation_buffers + .append_row_indices + .is_empty() + { + return Ok(()); + } + + let iter = self.group_values.iter_mut().zip(cols.iter()); + for (group_column, col) in iter { + group_column.vectorized_append( + col, + &self.vectorized_operation_buffers.append_row_indices, + )?; + } + + Ok(()) + } + + /// Perform `vectorized_equal_to` + /// + /// 1. Perform `vectorized_equal_to` for `rows` in `vectorized_equal_to_group_indices` + /// and `group_indices` in `vectorized_equal_to_group_indices`. + /// + /// 2. Check `equal_to_results`: + /// + /// If found equal to `rows`, set the `group_indices` to `rows` in `groups`. + /// + /// If found not equal to `row`s, just add them to `scalarized_indices`, + /// and perform `scalarized_intern` for them after. + /// Usually, such `rows` having same hash but different value with `exists rows` + /// are very few. + fn vectorized_equal_to(&mut self, cols: &[ArrayRef], groups: &mut [usize]) { + assert_eq!( + self.vectorized_operation_buffers + .equal_to_group_indices + .len(), + self.vectorized_operation_buffers.equal_to_row_indices.len() + ); + + self.vectorized_operation_buffers + .remaining_row_indices + .clear(); + + if self + .vectorized_operation_buffers + .equal_to_group_indices + .is_empty() + { + return; + } + + // 1. Perform `vectorized_equal_to` for `rows` in `vectorized_equal_to_group_indices` + // and `group_indices` in `vectorized_equal_to_group_indices` + let n = self + .vectorized_operation_buffers + .equal_to_group_indices + .len(); + let mut equal_to_results = mem::replace( + &mut self.vectorized_operation_buffers.equal_to_results, + BooleanBufferBuilder::new(0), + ); + equal_to_results.truncate(0); + equal_to_results.append_n(n, true); + + for (col_idx, group_col) in self.group_values.iter().enumerate() { + group_col.vectorized_equal_to( + &self.vectorized_operation_buffers.equal_to_group_indices, + &cols[col_idx], + &self.vectorized_operation_buffers.equal_to_row_indices, + &mut equal_to_results, + ); + } + + // 2. Check `equal_to_results`, if found not equal to `row`s, just add them + // to `scalarized_indices`, and perform `scalarized_intern` for them after. + let mut current_row_equal_to_result = false; + for (idx, &row) in self + .vectorized_operation_buffers + .equal_to_row_indices + .iter() + .enumerate() + { + let equal_to_result = equal_to_results.get_bit(idx); + + // Equal to case, set the `group_indices` to `rows` in `groups` + if equal_to_result { + groups[row] = + self.vectorized_operation_buffers.equal_to_group_indices[idx]; + } + current_row_equal_to_result |= equal_to_result; + + // Look forward next one row to check if have checked all results + // of current row + let next_row = self + .vectorized_operation_buffers + .equal_to_row_indices + .get(idx + 1) + .unwrap_or(&usize::MAX); + + // Have checked all results of current row, check the total result + if row != *next_row { + // Not equal to case, add `row` to `scalarized_indices` + if !current_row_equal_to_result { + self.vectorized_operation_buffers + .remaining_row_indices + .push(row); + } + + // Init the total result for checking next row + current_row_equal_to_result = false; + } + } + + self.vectorized_operation_buffers.equal_to_results = equal_to_results; + } + + /// It is possible that some `input rows` have the same + /// hash values with the `exist rows`, but have the different + /// actual values the exists. + /// + /// We can found them in `vectorized_equal_to`, and put them + /// into `scalarized_indices`. And for these `input rows`, + /// we will perform the `scalarized_intern` similar as what in + /// [`GroupValuesColumn`]. + /// + /// This design can make the process simple and still efficient enough: + /// + /// # About making the process simple + /// + /// Some corner cases become really easy to solve, like following cases: + /// + /// ```text + /// input row1 (same hash value with exist rows, but value different) + /// input row1 + /// ... + /// input row1 + /// ``` + /// + /// After performing `vectorized_equal_to`, we will found multiple `input rows` + /// not equal to the `exist rows`. However such `input rows` are repeated, only + /// one new group should be create for them. + /// + /// If we don't fallback to `scalarized_intern`, it is really hard for us to + /// distinguish the such `repeated rows` in `input rows`. And if we just fallback, + /// it is really easy to solve, and the performance is at least not worse than origin. + /// + /// # About performance + /// + /// The hash collision may be not frequent, so the fallback will indeed hardly happen. + /// In most situations, `scalarized_indices` will found to be empty after finishing to + /// perform `vectorized_equal_to`. + fn scalarized_intern_remaining( + &mut self, + cols: &[ArrayRef], + batch_hashes: &[u64], + groups: &mut [usize], + ) -> Result<()> { + if self + .vectorized_operation_buffers + .remaining_row_indices + .is_empty() + { + return Ok(()); + } + + let mut map = mem::take(&mut self.map); + + for &row in &self.vectorized_operation_buffers.remaining_row_indices { + let target_hash = batch_hashes[row]; + let entry = map.find_mut(target_hash, |(exist_hash, _)| { + // Somewhat surprisingly, this closure can be called even if the + // hash doesn't match, so check the hash first with an integer + // comparison first avoid the more expensive comparison with + // group value. https://github.com/apache/datafusion/pull/11718 + target_hash == *exist_hash + }); + + // Only `rows` having the same hash value with `exist rows` but different value + // will be process in `scalarized_intern`. + // So related `buckets` in `map` is ensured to be `Some`. + let Some((_, group_index_view)) = entry else { + unreachable!() + }; + + // Perform scalarized equal to + if self.scalarized_equal_to_remaining(group_index_view, cols, row, groups) { + // Found the row actually exists in group values, + // don't need to create new group for it. + continue; + } + + // Insert the `row` to `group_values` before checking `next row` + let group_idx = self.group_values[0].len(); + let mut checklen = 0; + for (i, group_value) in self.group_values.iter_mut().enumerate() { + group_value.append_val(&cols[i], row)?; + let len = group_value.len(); + if i == 0 { + checklen = len; + } else { + debug_assert_eq!(checklen, len); + } + } + + // Check if the `view` is `inlined` or `non-inlined` + if group_index_view.is_non_inlined() { + // Non-inlined case, get `group_index_list` from `group_index_lists`, + // then add the new `group` with the same hash values into it. + let list_offset = group_index_view.value() as usize; + let group_index_list = &mut self.group_index_lists[list_offset]; + group_index_list.push(group_idx); + } else { + // Inlined case + let list_offset = self.group_index_lists.len(); + + // Create new `group_index_list` including + // `exist group index` + `new group index`. + // Add new `group_index_list` into ``group_index_lists`. + let exist_group_index = group_index_view.value() as usize; + let new_group_index_list = vec![exist_group_index, group_idx]; + self.group_index_lists.push(new_group_index_list); + + // Update the `group_index_view` to non-inlined + let new_group_index_view = + GroupIndexView::new_non_inlined(list_offset as u64); + *group_index_view = new_group_index_view; + } + + groups[row] = group_idx; + } + + self.map = map; + Ok(()) + } + + fn scalarized_equal_to_remaining( + &self, + group_index_view: &GroupIndexView, + cols: &[ArrayRef], + row: usize, + groups: &mut [usize], + ) -> bool { + // Check if this row exists in `group_values` + fn check_row_equal( + array_row: &dyn GroupColumn, + lhs_row: usize, + array: &ArrayRef, + rhs_row: usize, + ) -> bool { + array_row.equal_to(lhs_row, array, rhs_row) + } + + if group_index_view.is_non_inlined() { + let list_offset = group_index_view.value() as usize; + let group_index_list = &self.group_index_lists[list_offset]; + + for &group_idx in group_index_list { + let mut check_result = true; + for (i, group_val) in self.group_values.iter().enumerate() { + if !check_row_equal(group_val.as_ref(), group_idx, &cols[i], row) { + check_result = false; + break; + } + } + + if check_result { + groups[row] = group_idx; + return true; + } + } + + // All groups unmatched, return false result + false + } else { + let group_idx = group_index_view.value() as usize; + for (i, group_val) in self.group_values.iter().enumerate() { + if !check_row_equal(group_val.as_ref(), group_idx, &cols[i], row) { + return false; + } + } + + groups[row] = group_idx; + true + } + } + + /// Return group indices of the hash, also if its `group_index_view` is non-inlined + #[cfg(test)] + fn get_indices_by_hash(&self, hash: u64) -> Option<(Vec<usize>, GroupIndexView)> { + let entry = self.map.find(hash, |(exist_hash, _)| hash == *exist_hash); + + match entry { + Some((_, group_index_view)) => { + if group_index_view.is_non_inlined() { + let list_offset = group_index_view.value() as usize; + Some(( + self.group_index_lists[list_offset].clone(), + *group_index_view, + )) + } else { + let group_index = group_index_view.value() as usize; + Some((vec![group_index], *group_index_view)) + } + } + None => None, + } + } +} + +/// instantiates a [`PrimitiveGroupValueBuilder`] and pushes it into $v +/// +/// Arguments: +/// `$v`: the vector to push the new builder into +/// `$nullable`: whether the input can contains nulls +/// `$t`: the primitive type of the builder +macro_rules! instantiate_primitive { + ($v:expr, $nullable:expr, $t:ty, $data_type:ident) => { + if $nullable { + let b = PrimitiveGroupValueBuilder::<$t, true>::new($data_type.to_owned()); + $v.push(Box::new(b) as _) + } else { + let b = PrimitiveGroupValueBuilder::<$t, false>::new($data_type.to_owned()); + $v.push(Box::new(b) as _) + } + }; +} + +/// Returns true if the specified data type has a specialized +/// [`GroupColumn`] builder in [`make_group_column`]. +/// +/// This is the allow-list that gates the `GroupValuesRows` fallback in +/// [`crate::aggregates::group_values::new_group_values`]: it must accept +/// exactly the set of types that [`make_group_column`] constructs a +/// builder for. The `group_column_supported_type_matches_make_group_column` +/// test below pins this biconditional. +fn group_column_supported_type(data_type: &DataType) -> bool { + // Nested types (Struct / List / LargeList / FixedSizeList, recursively) have + // no type-specialized `GroupColumn`; they are handled by the generic + // row-backed fallback in `make_group_column` whenever arrow's row format can + // encode them. Gate the fallback to nested types so intentionally-excluded + // scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the + // `group_column_supported_type` ⇔ `make_group_column` invariant holds. + if data_type.is_nested() { + return RowsGroupColumn::supports_type(data_type); + } + matches!( + *data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + // Only non-negative widths: a negative width is not a valid + // Arrow type (no array can be constructed for it), and the + // dispatcher in `make_group_column` rejects it. Keep the two + // in lockstep. + | DataType::FixedSizeBinary(0..) + | DataType::Date32 + | DataType::Date64 + // Only the semantically valid Time variants per the Arrow spec. + // The dispatcher in `make_group_column` returns NotImpl for the + // other unit combinations, so accepting them here would cause a + // schema to be routed into GroupValuesColumn and then fail at + // intern. Keep these two arms in lockstep with the dispatcher. + | DataType::Time32(TimeUnit::Second) + | DataType::Time32(TimeUnit::Millisecond) + | DataType::Time64(TimeUnit::Microsecond) + | DataType::Time64(TimeUnit::Nanosecond) + | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Interval(_) + | DataType::Utf8View + | DataType::BinaryView + | DataType::Boolean + ) +} + +/// Build a [`GroupColumn`] for a single schema field. +/// +/// Extracted from the inline match that used to live in +/// [`GroupValuesColumn::intern`] so the per-field dispatch lives in one +/// place. This factory is the single source of truth for which Arrow types +/// map to which builder, and it is the function that future nested-type +/// specializations (e.g. `Struct`, `List`, `LargeList`) plug into without +/// having to enumerate every combination inline. +/// +/// Returns `Err(not_impl_err!(...))` for any type not in the supported set; +/// callers (`GroupValues::intern`) propagate that error so the +/// `GroupValuesRows` fallback can take over upstream of this builder. +/// +/// The allow-list that gates this dispatcher lives in +/// [`group_column_supported_type`] directly above. +fn make_group_column(field: &Field) -> Result<Box<dyn GroupColumn>> { + let nullable = field.is_nullable(); + let data_type = field.data_type(); + let mut v: Vec<Box<dyn GroupColumn>> = Vec::with_capacity(1); + match *data_type { + DataType::Int8 => instantiate_primitive!(v, nullable, Int8Type, data_type), + DataType::Int16 => instantiate_primitive!(v, nullable, Int16Type, data_type), + DataType::Int32 => instantiate_primitive!(v, nullable, Int32Type, data_type), + DataType::Int64 => instantiate_primitive!(v, nullable, Int64Type, data_type), + DataType::UInt8 => instantiate_primitive!(v, nullable, UInt8Type, data_type), + DataType::UInt16 => instantiate_primitive!(v, nullable, UInt16Type, data_type), + DataType::UInt32 => instantiate_primitive!(v, nullable, UInt32Type, data_type), + DataType::UInt64 => instantiate_primitive!(v, nullable, UInt64Type, data_type), + DataType::Float16 => { + instantiate_primitive!(v, nullable, Float16Type, data_type) + } + DataType::Float32 => { + instantiate_primitive!(v, nullable, Float32Type, data_type) + } + DataType::Float64 => { + instantiate_primitive!(v, nullable, Float64Type, data_type) + } + DataType::Date32 => instantiate_primitive!(v, nullable, Date32Type, data_type), + DataType::Date64 => instantiate_primitive!(v, nullable, Date64Type, data_type), + DataType::Time32(t) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, Time32SecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, Time32MillisecondType, data_type) + } + // Time32 with Microsecond / Nanosecond is not a valid Arrow type + // combination; reject explicitly so group_column_supported_type + // and this dispatcher stay in lockstep (see consistency fuzz below). + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + }, + DataType::Time64(t) => match t { + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, Time64MicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, Time64NanosecondType, data_type) + } + // Time64 with Second / Millisecond is not a valid Arrow type + // combination; reject explicitly. + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + }, + DataType::Timestamp(t, _) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, TimestampSecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, TimestampMillisecondType, data_type) + } + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, TimestampMicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, TimestampNanosecondType, data_type) + } + }, + DataType::Duration(t) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, DurationSecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, DurationMillisecondType, data_type) + } + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, DurationMicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, DurationNanosecondType, data_type) + } + }, + // `IntervalUnit` has exactly three variants, so this match is exhaustive + // with no fallback arm (unlike Time32 / Time64). + DataType::Interval(u) => match u { + IntervalUnit::YearMonth => { + instantiate_primitive!(v, nullable, IntervalYearMonthType, data_type) + } + IntervalUnit::DayTime => { + instantiate_primitive!(v, nullable, IntervalDayTimeType, data_type) + } + IntervalUnit::MonthDayNano => { + instantiate_primitive!(v, nullable, IntervalMonthDayNanoType, data_type) + } + }, + DataType::Decimal128(_, _) => { + instantiate_primitive!(v, nullable, Decimal128Type, data_type) + } + DataType::Decimal256(_, _) => { + instantiate_primitive!(v, nullable, Decimal256Type, data_type) + } + DataType::Utf8 => { + v.push(Box::new(ByteGroupValueBuilder::<i32>::new( + OutputType::Utf8, + ))); + } + DataType::LargeUtf8 => { + v.push(Box::new(ByteGroupValueBuilder::<i64>::new( + OutputType::Utf8, + ))); + } + DataType::Binary => { + v.push(Box::new(ByteGroupValueBuilder::<i32>::new( + OutputType::Binary, + ))); + } + DataType::LargeBinary => { + v.push(Box::new(ByteGroupValueBuilder::<i64>::new( + OutputType::Binary, + ))); + } + // A negative width is not a valid Arrow type; it falls to the `_` + // arm below, matching `group_column_supported_type`. + DataType::FixedSizeBinary(byte_width @ 0..) => { + v.push(Box::new(FixedSizeBinaryGroupValueBuilder::new(byte_width))); + } + DataType::Utf8View => { + v.push(Box::new(ByteViewGroupValueBuilder::<StringViewType>::new())); + } + DataType::BinaryView => { + v.push(Box::new(ByteViewGroupValueBuilder::<BinaryViewType>::new())); + } + DataType::Boolean => { + if nullable { + v.push(Box::new(BooleanGroupValueBuilder::<true>::new())); + } else { + v.push(Box::new(BooleanGroupValueBuilder::<false>::new())); + } + } + // Generic fallback for nested types (Struct / List / LargeList / + // FixedSizeList, recursively) that lack a type-specialized builder but + // can be encoded by arrow's row format. This is what lets a mixed + // schema keep the column-wise fast path for its native columns instead + // of dropping the whole key onto `GroupValuesRows`. + ref dt if dt.is_nested() && RowsGroupColumn::supports_type(dt) => { + v.push(Box::new(RowsGroupColumn::try_new(dt.clone())?)); + } + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + } + debug_assert_eq!( + v.len(), + 1, + "make_group_column must push exactly one builder" + ); + Ok(v.into_iter().next().unwrap()) +} + +impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + // `try_new` and the reset points in `emit` / `clear_shrink` keep + // `self.group_values` populated with one builder per schema field, + // so no lazy initialization is needed here. + if !STREAMING { + self.vectorized_intern(cols, groups) + } else { + self.scalarized_intern(cols, groups) + } + } + + fn size(&self) -> usize { + let group_values_size: usize = self.group_values.iter().map(|v| v.size()).sum(); + group_values_size + self.map_size + self.hashes_buffer.allocated_size() + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { + if self.group_values.is_empty() { + return 0; + } + + self.group_values[0].len() + } + + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + let mut output = match emit_to { + EmitTo::All => { + // Replace the column builders with a fresh set so the + // aggregator is immediately reusable after the drain. + // Same `self.schema` was already validated by `try_new`, + // so `build_group_columns` would only error here if some + // out-of-band schema mutation occurred — propagate it as + // a real Result rather than panicking. + let fresh = Self::build_group_columns(&self.schema)?; + let group_values = mem::replace(&mut self.group_values, fresh); + + group_values + .into_iter() + .map(|v| v.build()) + .collect::<Vec<_>>() + } + EmitTo::First(n) => { + let output = self + .group_values + .iter_mut() + .map(|v| v.take_n(n)) + .collect::<Vec<_>>(); + let mut next_new_list_offset = 0; + + self.map.retain(|(_exist_hash, group_idx_view)| { + // In non-streaming case, we need to check if the `group index view` + // is `inlined` or `non-inlined` + if !STREAMING && group_idx_view.is_non_inlined() { + // Non-inlined case + // We take `group_index_list` from `old_group_index_lists` + + // list_offset is incrementally + self.emit_group_index_list_buffer.clear(); + let list_offset = group_idx_view.value() as usize; + for group_index in self.group_index_lists[list_offset].iter() { + if let Some(remaining) = group_index.checked_sub(n) { + self.emit_group_index_list_buffer.push(remaining); + } + } + + // The possible results: + // - `new_group_index_list` is empty, we should erase this bucket + // - only one value in `new_group_index_list`, switch the `view` to `inlined` + // - still multiple values in `new_group_index_list`, build and set the new `unlined view` + if self.emit_group_index_list_buffer.is_empty() { + false + } else if self.emit_group_index_list_buffer.len() == 1 { + let group_index = + self.emit_group_index_list_buffer.first().unwrap(); + *group_idx_view = + GroupIndexView::new_inlined(*group_index as u64); + true + } else { + let group_index_list = + &mut self.group_index_lists[next_new_list_offset]; + group_index_list.clear(); + group_index_list + .extend(self.emit_group_index_list_buffer.iter()); + *group_idx_view = GroupIndexView::new_non_inlined( + next_new_list_offset as u64, + ); + next_new_list_offset += 1; + true + } + } else { + // In `streaming case`, the `group index view` is ensured to be `inlined` + debug_assert!(!group_idx_view.is_non_inlined()); + + // Inlined case, we just decrement group index by n) + let group_index = group_idx_view.value() as usize; + match group_index.checked_sub(n) { + // Group index was >= n, shift value down + Some(sub) => { + *group_idx_view = GroupIndexView::new_inlined(sub as u64); + true + } + // Group index was < n, so remove from table + None => false, + } + } + }); + + if !STREAMING { + self.group_index_lists.truncate(next_new_list_offset); + } + + output + } + }; + + // TODO: Materialize dictionaries in group keys (#7647) + for (field, array) in self.schema.fields.iter().zip(&mut output) { + let expected = field.data_type(); + if let DataType::Dictionary(_, v) = expected { + let actual = array.data_type(); + if v.as_ref() != actual { + return Err(internal_datafusion_err!( + "Converted group rows expected dictionary of {v} got {actual}" + )); + } + *array = cast(array.as_ref(), expected)?; + } + } + + Ok(output) + } + + fn clear_shrink(&mut self, num_rows: usize) { + // Reset to a fresh column-builder vector. The schema was validated + // in `try_new`, so rebuilding cannot fail unless something else + // mutated the schema out-of-band — surface that as a panic since + // `clear_shrink` is infallible by trait signature. + self.group_values = Self::build_group_columns(&self.schema) + .expect("schema previously validated in try_new"); + self.map.clear(); + self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); + self.hashes_buffer.clear(); + self.hashes_buffer.shrink_to(num_rows); + + // Such structures are only used in `non-streaming` case + if !STREAMING { + self.group_index_lists.clear(); + self.emit_group_index_list_buffer.clear(); + self.vectorized_operation_buffers.clear(); + } + } +} + +/// Returns true if [`GroupValuesColumn`] supported for the specified schema +pub fn supported_schema(schema: &Schema) -> bool { + schema + .fields() + .iter() + .map(|f| f.data_type()) + .all(group_column_supported_type) +} + +///Shows how many `null`s there are in an array +enum Nulls { + /// All array items are `null`s + All, + /// There are both `null`s and non-`null`s in the array items + Some, + /// There are no `null`s in the array items + None, +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::Arc}; + + use arrow::array::{ + Array, ArrayRef, DurationMicrosecondArray, FixedSizeBinaryArray, Float16Array, + Int32Array, Int64Array, PrimitiveArray, RecordBatch, StringArray, + StringViewArray, + }; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; + use datafusion_common::utils::proxy::HashTableAllocExt; + use datafusion_expr::EmitTo; + + use crate::aggregates::group_values::{ + GroupValues, multi_group_by::GroupValuesColumn, + }; + + use super::{ + GroupIndexView, group_column_supported_type, make_group_column, supported_schema, + }; + + /// A mixed group-by key of several native columns plus one nested column + /// that has no type-specialized `GroupColumn`. + /// + /// Before the generic row-backed fallback, `supported_schema` returned + /// `false` for this schema, so the *entire* key dropped to the row-wise + /// `GroupValuesRows`. Now only the nested column pays the row-encoding + /// cost; the native columns keep their compact column-wise storage. This + /// test proves both that (a) the results are identical and (b) the + /// column-wise path now uses less memory than the all-rows fallback. + #[test] + fn mixed_schema_column_path_uses_less_memory_than_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Int64Array}; + use arrow::datatypes::Int64Type; + + // 8 native Int64 columns + 1 FixedSizeList<Int64, 4> ("embedding"). + let fsl_field = Arc::new(Field::new("item", DataType::Int64, true)); + let mut fields: Vec<Field> = (0..8) + .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) + .collect(); + fields.push(Field::new( + "emb", + DataType::FixedSizeList(Arc::clone(&fsl_field), 4), + true, + )); + let schema: SchemaRef = Arc::new(Schema::new(fields)); + + // The whole schema must now be eligible for the column-wise path. + assert!( + supported_schema(schema.as_ref()), + "mixed native + nested schema should be column-supported now" + ); + + // Build `n_groups` distinct rows (each row is its own group). + let n_groups = 4000usize; + let mut cols: Vec<ArrayRef> = (0..8) + .map(|c| { + let vals: Vec<i64> = + (0..n_groups).map(|r| (r as i64) * 8 + c as i64).collect(); + Arc::new(Int64Array::from(vals)) as ArrayRef + }) + .collect(); + let emb: Vec<Option<Vec<Option<i64>>>> = (0..n_groups) + .map(|r| { + Some(vec![ + Some(r as i64), + Some(r as i64 + 1), + Some(r as i64 + 2), + Some(r as i64 + 3), + ]) + }) + .collect(); + cols.push( + Arc::new(FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>( + emb, 4, + )) as ArrayRef, + ); + + // Intern the same data into both implementations. + let mut column_path = GroupValuesColumn::<false>::try_new(Arc::clone(&schema)) + .expect("column path"); + let mut rows_path = + GroupValuesRows::try_new(Arc::clone(&schema)).expect("rows path"); + + let mut g1 = vec![]; + let mut g2 = vec![]; + column_path.intern(&cols, &mut g1).unwrap(); + rows_path.intern(&cols, &mut g2).unwrap(); + + // (a) Correctness: same number of groups and identical group assignment. + assert_eq!(column_path.len(), n_groups); + assert_eq!(rows_path.len(), n_groups); + assert_eq!(g1, g2, "group assignment must match the rows fallback"); + + // (b) Memory: the column-wise path stores the 8 native columns compactly + // and only row-encodes the nested one, so it should be smaller than + // encoding every column into rows. + // + // The delta is only printed here — a hard `column_size < rows_size` + // assert would be brittle to future Arrow row-format or memory- + // accounting changes without reflecting a grouping-correctness + // regression. Track the memory improvement via benchmarks instead. + let column_size = column_path.size(); + let rows_size = rows_path.size(); + println!( + "mixed-schema group values size: column-wise = {column_size} bytes, \ + all-rows fallback = {rows_size} bytes \ + ({:.1}% of fallback)", + 100.0 * column_size as f64 / rows_size as f64 + ); + + // Emitted values must be equal too (compare via the rows fallback which + // is the established reference implementation). + let out_col = column_path.emit(EmitTo::All).unwrap(); + let out_row = rows_path.emit(EmitTo::All).unwrap(); + assert_eq!(out_col.len(), out_row.len()); + for (a, b) in out_col.iter().zip(out_row.iter()) { + assert_eq!(a.as_ref(), b.as_ref()); + } + } + + /// Relabel a group-index vector so labels are assigned in order of first + /// appearance. Two vectors are equivalent groupings iff their canonical + /// forms are equal — this ignores the (opaque, non-semantic) difference in + /// group-index numbering between the vectorized column path and the + /// sequential rows fallback. + /// + /// The [`GroupValues`] trait only guarantees that equal keys receive the + /// same group-id and that new keys receive a fresh id; the order in which + /// new ids are handed out is deliberately not part of the contract, and + /// can differ between correct implementations (e.g. because of internal + /// hash-map ordering). Canonicalizing before comparison is what lets us + /// assert equivalence across implementations. + fn canonical_grouping(groups: &[usize]) -> Vec<usize> { + let mut map = HashMap::new(); + let mut next = 0usize; + groups + .iter() + .map(|&g| { + *map.entry(g).or_insert_with(|| { + let v = next; + next += 1; + v + }) + }) + .collect() + } + + /// The generic row-backed column must be behavior-preserving: for the + /// nested columns it now handles, `GroupValuesColumn` must induce the same + /// grouping (partition of rows) as the established `GroupValuesRows` + /// fallback — including the float `-0.0` / `+0.0` / `NaN` edge cases decided + /// jointly by hashing and the row format. + #[test] + fn nested_float_edge_cases_match_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Float64Array}; + + let item = Arc::new(Field::new("item", DataType::Float64, true)); + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "emb", + DataType::FixedSizeList(Arc::clone(&item), 2), + true, + )])); + assert!(supported_schema(schema.as_ref())); + + // Rows exercising +0.0 vs -0.0, two NaN bit patterns, and inner nulls. + let nan = f64::NAN; + let other_nan = f64::from_bits(0x7ff8_0000_0000_0001); + let values = Float64Array::from(vec![ + Some(0.0), + Some(1.0), // [ +0.0, 1.0 ] + Some(-0.0), + Some(1.0), // [ -0.0, 1.0 ] + Some(nan), + Some(2.0), // [ NaN, 2.0 ] + Some(other_nan), + Some(2.0), // [ NaN', 2.0 ] + Some(0.0), + Some(1.0), // [ +0.0, 1.0 ] (dup of row 0) + ]); + let field_ref = Arc::new(Field::new("item", DataType::Float64, true)); + let input: ArrayRef = Arc::new(FixedSizeListArray::new( + field_ref, + 2, + Arc::new(values), + None, + )); + + let cols = vec![input]; + + let mut column_path = + GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap(); + let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap(); + + let mut g1 = vec![]; + let mut g2 = vec![]; + column_path.intern(&cols, &mut g1).unwrap(); + rows_path.intern(&cols, &mut g2).unwrap(); + + assert_eq!( + canonical_grouping(&g1), + canonical_grouping(&g2), + "column-wise path must induce the same grouping as the rows fallback \ + on float edge cases (got column={g1:?}, rows={g2:?})" + ); + assert_eq!(column_path.len(), rows_path.len()); + } + + /// Equivalence across multiple `intern` batches and `EmitTo::First(n)`. + #[test] + fn multi_batch_and_emit_first_matches_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Int32Array}; + use arrow::datatypes::Int32Type; + + let item = Arc::new(Field::new("item", DataType::Int32, true)); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("emb", DataType::FixedSizeList(Arc::clone(&item), 2), true), + ])); + + let make_batch = |base: i32| -> Vec<ArrayRef> { + let k = Arc::new(Int32Array::from(vec![base, base + 1, base])) as ArrayRef; + let emb: Vec<Option<Vec<Option<i32>>>> = vec![ + Some(vec![Some(base), Some(base)]), + Some(vec![Some(base + 1), None]), + Some(vec![Some(base), Some(base)]), // dup of row 0 + ]; + let emb = Arc::new( + FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>(emb, 2), + ) as ArrayRef; + vec![k, emb] + }; + + let mut column_path = + GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap(); + let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap(); + + for base in [0, 10, 0] { + let cols = make_batch(base); + let (mut a, mut b) = (vec![], vec![]); + column_path.intern(&cols, &mut a).unwrap(); + rows_path.intern(&cols, &mut b).unwrap(); + // Same grouping (partition), even if the opaque group-index labels + // differ between the vectorized and sequential paths. + assert_eq!( + canonical_grouping(&a), + canonical_grouping(&b), + "grouping must match for batch base={base}" + ); + } + + let total_groups = column_path.len(); + assert_eq!(total_groups, rows_path.len()); + + // `EmitTo::First(n)` then `EmitTo::All` on the nested column path must + // work and together emit exactly `total_groups` rows. (Cross-path value + // equality is covered by `mixed_schema_...` and the row_backed unit + // tests; group-index ordering differs here so we check counts.) + let col_first = column_path.emit(EmitTo::First(2)).unwrap(); + assert_eq!(col_first[0].len(), 2); + let col_rest = column_path.emit(EmitTo::All).unwrap(); + assert_eq!(col_first[0].len() + col_rest[0].len(), total_groups); + // Column count / schema preserved on both emits. + assert_eq!(col_first.len(), schema.fields().len()); + assert_eq!(col_rest.len(), schema.fields().len()); + } + + /// CRITICAL invariant: if `group_column_supported_type(t)` returns true + /// the dispatcher must accept that type at intern time, and conversely + /// if `group_column_supported_type(t)` returns false the planner must + /// NOT route it through `GroupValuesColumn`. A divergence here would + /// let the planner select `GroupValuesColumn` for a type whose + /// dispatcher arm is missing, producing a runtime `not_impl_err` after + /// the field reaches the builder factory. + /// + /// This test fuzzes a representative cross-section of types and asserts + /// both directions of the biconditional. When a new specialization is + /// added (`Float16`, `FixedSizeList`, `Struct`, ...) it should be added + /// to the supported_cases vector; when a type is intentionally rejected + /// it should be added to unsupported_cases. + #[test] + fn group_column_supported_type_matches_make_group_column() { + let supported_cases: Vec<DataType> = vec![ + DataType::Int8, + DataType::Int64, + DataType::UInt64, + DataType::Float32, + DataType::Float64, + DataType::Float16, + DataType::Decimal128(38, 10), + DataType::Decimal256(76, 10), + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::FixedSizeBinary(16), + // Zero-width FixedSizeBinary is valid per the Arrow spec + DataType::FixedSizeBinary(0), + DataType::Boolean, + DataType::Date32, + DataType::Date64, + DataType::Time32(arrow::datatypes::TimeUnit::Second), + DataType::Time32(arrow::datatypes::TimeUnit::Millisecond), + DataType::Time64(arrow::datatypes::TimeUnit::Microsecond), + DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + DataType::Duration(arrow::datatypes::TimeUnit::Second), + DataType::Duration(arrow::datatypes::TimeUnit::Millisecond), + DataType::Duration(arrow::datatypes::TimeUnit::Microsecond), + DataType::Duration(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth), + DataType::Interval(arrow::datatypes::IntervalUnit::DayTime), + DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano), + ]; + + for dt in &supported_cases { + assert!( + group_column_supported_type(dt), + "expected group_column_supported_type=true for {dt:?}" + ); + let field = Field::new("col", dt.clone(), true); + make_group_column(&field).unwrap_or_else(|e| { + panic!( + "group_column_supported_type accepted {dt:?} but make_group_column rejected: {e}" + ) + }); + } + + let unsupported_cases: Vec<DataType> = vec![ + // Invalid Time-unit combinations: Time32 is defined only for + // Second / Millisecond and Time64 only for Microsecond / + // Nanosecond. The TimeUnit enum allows constructing the other + // combinations programmatically, but they are not valid Arrow + // types and must be rejected by both group_column_supported_type + // and the dispatcher. + DataType::Time64(arrow::datatypes::TimeUnit::Second), + DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), + DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), + DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), + // A negative width is representable in the DataType but is not + // a valid Arrow type; no array can be constructed for it. + DataType::FixedSizeBinary(-5), + ]; + + for dt in &unsupported_cases { + assert!( + !group_column_supported_type(dt), + "expected group_column_supported_type=false for {dt:?}" + ); + let field = Field::new("col", dt.clone(), true); + assert!( + make_group_column(&field).is_err(), + "group_column_supported_type rejected {dt:?} but make_group_column accepted it" + ); + } + } + + // `Duration` group keys stay on the `GroupValuesColumn` fast path, dedup + // (including nulls), and round-trip with the `Duration` type preserved. + #[test] + fn test_group_values_column_duration() { + use arrow::datatypes::TimeUnit; + + let schema = Arc::new(Schema::new(vec![ + Field::new("d", DataType::Duration(TimeUnit::Microsecond), true), + Field::new("i", DataType::Int64, true), + ])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap(); + + // (d, i) rows, where row 3 repeats row 0 and row 4 repeats the null pair. + let d: ArrayRef = Arc::new(DurationMicrosecondArray::from(vec![ + Some(10), + None, + Some(20), + Some(10), + None, + ])); + let i: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + Some(1), + None, + ])); + let mut groups = Vec::new(); + group_values.intern(&[d, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0, 1]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // The Duration column round-trips as Duration on emit, not bare i64. + assert_eq!( + emitted[0].data_type(), + &DataType::Duration(TimeUnit::Microsecond) + ); + let actual = emitted[0] + .as_any() + .downcast_ref::<DurationMicrosecondArray>() + .expect("emitted column should be a DurationMicrosecondArray"); + // Three groups in first-seen order: 10, null, 20. + assert_eq!(actual.len(), 3); + assert_eq!(actual.value(0), 10); + assert!(actual.is_null(1)); + assert_eq!(actual.value(2), 20); + } + + // `(Float16, Int32)` keys: ±0.0 collapse (stored as +0.0), NaNs collapse, and + // the Int32 key keeps `(0.0, 4)` distinct from `(±0.0, 3)`. + #[test] + fn test_group_values_column_float16() { + use half::f16; + + let schema = Arc::new(Schema::new(vec![ + Field::new("f", DataType::Float16, true), + Field::new("i", DataType::Int32, true), + ])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap(); + + let f: ArrayRef = Arc::new(Float16Array::from(vec![ + Some(f16::from_f32(1.0)), + Some(f16::from_f32(-0.0)), + Some(f16::from_f32(0.0)), + Some(f16::from_f32(0.0)), + Some(f16::NAN), + Some(f16::NAN), + None, + None, + ])); + let i: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(3), + Some(3), + Some(3), + Some(4), + Some(3), + Some(3), + Some(3), + Some(3), + ])); + let mut groups = Vec::new(); + group_values.intern(&[f, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 1, 2, 3, 3, 4, 4]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + assert_eq!(emitted[0].data_type(), &DataType::Float16); + let keys = emitted[0] + .as_any() + .downcast_ref::<Float16Array>() + .expect("emitted column should be a Float16Array"); + assert_eq!(keys.len(), 5); + assert_eq!(keys.value(0), f16::from_f32(1.0)); + // The ±0.0 group is stored canonically as +0.0 (not -0.0). + assert_eq!(keys.value(1).to_bits(), f16::from_f32(0.0).to_bits()); + assert_eq!(keys.value(2).to_bits(), f16::from_f32(0.0).to_bits()); + assert!(keys.value(3).is_nan()); + assert!(keys.is_null(4)); + let ids = emitted[1] + .as_any() + .downcast_ref::<Int32Array>() + .expect("emitted column should be an Int32Array"); + assert_eq!(ids.values().to_vec(), vec![3, 3, 4, 3, 3]); + } + + // `(Interval, Int32)` keys for each of the three interval units: null keys + // dedup, the Int32 key splits equal intervals, and emit gives back Interval. + #[test] + fn test_group_values_column_interval() { + use arrow::datatypes::{ + ArrowPrimitiveType, IntervalDayTime, IntervalDayTimeType, + IntervalMonthDayNano, IntervalMonthDayNanoType, IntervalUnit, + IntervalYearMonthType, + }; + + fn check<T: ArrowPrimitiveType>(unit: IntervalUnit, value: T::Native) { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Interval(unit), true), + Field::new("n", DataType::Int32, true), + ])); + assert!(supported_schema(&schema), "{unit:?} schema not supported"); + let mut group_values = + GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap(); + + let i: ArrayRef = Arc::new(PrimitiveArray::<T>::from_iter([ + Some(value), + None, + Some(value), + None, + Some(value), + ])); + let n: ArrayRef = Arc::new(Int32Array::from(vec![3, 3, 3, 3, 4])); + let mut groups = Vec::new(); + group_values.intern(&[i, n], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 0, 1, 2], "{unit:?}"); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // The emitted key keeps its Interval type, not the bare native. + assert_eq!(emitted[0].data_type(), &DataType::Interval(unit)); + let actual = emitted[0] + .as_any() + .downcast_ref::<PrimitiveArray<T>>() + .unwrap_or_else(|| panic!("emitted column should be a {unit:?} array")); + // Three groups in first-seen order: value, null, value (n=4). + assert_eq!(actual.len(), 3, "{unit:?}"); + assert_eq!(actual.value(0), value, "{unit:?}"); + assert!(actual.is_null(1), "{unit:?}"); + assert_eq!(actual.value(2), value, "{unit:?}"); + let ids = emitted[1] + .as_any() + .downcast_ref::<Int32Array>() + .expect("emitted column should be an Int32Array"); + assert_eq!(ids.values().to_vec(), vec![3, 3, 4], "{unit:?}"); + } + + check::<IntervalYearMonthType>(IntervalUnit::YearMonth, 13); + check::<IntervalDayTimeType>(IntervalUnit::DayTime, IntervalDayTime::new(1, 500)); + check::<IntervalMonthDayNanoType>( + IntervalUnit::MonthDayNano, + IntervalMonthDayNano::new(1, 0, 0), + ); + } + + #[test] + fn supported_schema_rejects_mix_of_supported_and_unsupported() { + // One unsupported column flips the whole schema to the GroupValuesRows + // fallback. Time64(Second) stays invalid as new primitive builders land. + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new( + "c", + DataType::Time64(arrow::datatypes::TimeUnit::Second), + true, + ), + ]); + assert!(!supported_schema(&schema)); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Boolean, true), + ]); + assert!(supported_schema(&schema)); + } + + #[test] + fn try_new_returns_not_impl_for_unsupported_top_level_type() { + // `try_new` now eagerly constructs the per-field GroupColumn + // builders via `make_group_column`, so an unsupported schema is + // rejected at construction time rather than at first `intern`. + // `GroupValuesColumn` doesn't implement `Debug`, so explicit match + // instead of `unwrap_err`. + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + DataType::Time64(arrow::datatypes::TimeUnit::Second), + true, + )])); + match GroupValuesColumn::<false>::try_new(schema) { + Ok(_) => panic!("expected NotImpl error, but try_new succeeded"), + Err(e) => { + let msg = e.to_string(); + assert!( + msg.contains("not supported in GroupValuesColumn"), + "expected NotImpl error from dispatcher, got: {msg}" + ); + } + } + } + + #[test] + fn test_intern_for_vectorized_group_values() { + let data_set = VectorizedTestDataSet::new(); + let mut group_values = + GroupValuesColumn::<false>::try_new(data_set.schema()).unwrap(); + + data_set.load_to_group_values(&mut group_values); + let actual_batch = group_values.emit(EmitTo::All).unwrap(); + let actual_batch = RecordBatch::try_new(data_set.schema(), actual_batch).unwrap(); + + check_result(&actual_batch, &data_set.expected_batch); + } + + #[test] + fn test_intern_for_fixed_size_binary_group_values() { + // Two-column group by `(FixedSizeBinary(2), Int64)` exercising the + // vectorized intern path end-to-end (hashing included), with nulls, + // within-batch repeats and across-batch repeats. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(2), true), + Field::new("b", DataType::Int64, true), + ])); + let mut group_values = + GroupValuesColumn::<false>::try_new(Arc::clone(&schema)).unwrap(); + + fn fsb(values: Vec<Option<&[u8; 2]>>) -> ArrayRef { + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.into_iter(), + 2, + ) + .unwrap(), + ) + } + + let batch1: Vec<ArrayRef> = vec![ + fsb(vec![Some(b"aa"), Some(b"aa"), None, None, Some(b"bb")]), + Arc::new(Int64Array::from(vec![ + Some(1), + Some(1), + None, + Some(2), + None, + ])), + ]; + // Mix of groups repeated from batch1 and new groups + let batch2: Vec<ArrayRef> = vec![ + fsb(vec![Some(b"aa"), Some(b"cc"), None, Some(b"bb")]), + Arc::new(Int64Array::from(vec![Some(1), Some(1), None, Some(3)])), + ]; + + group_values.intern(&batch1, &mut vec![]).unwrap(); + group_values.intern(&batch2, &mut vec![]).unwrap(); + + let actual_batch = group_values.emit(EmitTo::All).unwrap(); + let actual_batch = + RecordBatch::try_new(Arc::clone(&schema), actual_batch).unwrap(); + + let expected_batch = RecordBatch::try_new( + schema, + vec![ + fsb(vec![ + Some(b"aa"), + None, + None, + Some(b"bb"), + Some(b"cc"), + Some(b"bb"), + ]), + Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(1), + Some(3), + ])), + ], + ) + .unwrap(); + + assert_eq!(actual_batch.num_rows(), expected_batch.num_rows()); + check_result(&actual_batch, &expected_batch); + } + + #[test] + fn test_emit_first_n_for_vectorized_group_values() { + let data_set = VectorizedTestDataSet::new(); + let mut group_values = + GroupValuesColumn::<false>::try_new(data_set.schema()).unwrap(); + + // 1~num_rows times to emit the groups + let num_rows = data_set.expected_batch.num_rows(); + let schema = data_set.schema(); + for times_to_take in 1..=num_rows { + // Write data after emitting + data_set.load_to_group_values(&mut group_values); + + // Emit `times_to_take` times, collect and concat the sub-results to total result, + // then check it + let suggest_num_emit = data_set.expected_batch.num_rows() / times_to_take; + let mut num_remaining_rows = num_rows; + let mut actual_sub_batches = Vec::new(); + + for nth_time in 0..times_to_take { + let num_emit = if nth_time == times_to_take - 1 { + num_remaining_rows + } else { + suggest_num_emit + }; + + let sub_batch = group_values.emit(EmitTo::First(num_emit)).unwrap(); + let sub_batch = + RecordBatch::try_new(Arc::clone(&schema), sub_batch).unwrap(); + actual_sub_batches.push(sub_batch); + + num_remaining_rows -= num_emit; + } + assert!(num_remaining_rows == 0); + + let actual_batch = concat_batches(&schema, &actual_sub_batches).unwrap(); + check_result(&actual_batch, &data_set.expected_batch); + } + } + + #[test] + fn test_hashtable_modifying_in_emit_first_n() { + // Situations should be covered: + // 1. Erase inlined group index view + // 2. Erase whole non-inlined group index view + // 3. Erase + decrease group indices in non-inlined group index view + // + view still non-inlined after decreasing + // 4. Erase + decrease group indices in non-inlined group index view + // + view switch to inlined after decreasing + // 5. Only decrease group index in inlined group index view + // 6. Only decrease group indices in non-inlined group index view + // 7. Erase all things + + let field = Field::new_list_field(DataType::Int32, true); + let schema = Arc::new(Schema::new_with_metadata(vec![field], HashMap::new())); + let mut group_values = GroupValuesColumn::<false>::try_new(schema).unwrap(); + + // Seed the column with 12 placeholder rows so the upcoming + // `emit(EmitTo::First(4))` calls can `take_n` without panicking. + // The hashmap entries below reference group indices 0..=11, so the + // single column builder needs at least 12 rows to back them. + let seed: ArrayRef = Arc::new(Int32Array::from(vec![0_i32; 12])); + for row in 0..12 { + group_values.group_values[0] + .append_val(&seed, row) + .expect("seed append"); + } + + // Insert group index views and check if success to insert + insert_inline_group_index_view(&mut group_values, 0, 0); + insert_non_inline_group_index_view(&mut group_values, 1, vec![1, 2]); + insert_non_inline_group_index_view(&mut group_values, 2, vec![3, 4, 5]); + insert_inline_group_index_view(&mut group_values, 3, 6); + insert_non_inline_group_index_view(&mut group_values, 4, vec![7, 8]); + insert_non_inline_group_index_view(&mut group_values, 5, vec![9, 10, 11]); + + assert_eq!( + group_values.get_indices_by_hash(0).unwrap(), + (vec![0], GroupIndexView::new_inlined(0)) + ); + assert_eq!( + group_values.get_indices_by_hash(1).unwrap(), + (vec![1, 2], GroupIndexView::new_non_inlined(0)) + ); + assert_eq!( + group_values.get_indices_by_hash(2).unwrap(), + (vec![3, 4, 5], GroupIndexView::new_non_inlined(1)) + ); + assert_eq!( + group_values.get_indices_by_hash(3).unwrap(), + (vec![6], GroupIndexView::new_inlined(6)) + ); + assert_eq!( + group_values.get_indices_by_hash(4).unwrap(), + (vec![7, 8], GroupIndexView::new_non_inlined(2)) + ); + assert_eq!( + group_values.get_indices_by_hash(5).unwrap(), + (vec![9, 10, 11], GroupIndexView::new_non_inlined(3)) + ); + assert_eq!(group_values.map.len(), 6); + + // Emit first 4 to test cases 1~3, 5~6 + let _ = group_values.emit(EmitTo::First(4)).unwrap(); + assert!(group_values.get_indices_by_hash(0).is_none()); + assert!(group_values.get_indices_by_hash(1).is_none()); + assert_eq!( + group_values.get_indices_by_hash(2).unwrap(), + (vec![0, 1], GroupIndexView::new_non_inlined(0)) + ); + assert_eq!( + group_values.get_indices_by_hash(3).unwrap(), + (vec![2], GroupIndexView::new_inlined(2)) + ); + assert_eq!( + group_values.get_indices_by_hash(4).unwrap(), + (vec![3, 4], GroupIndexView::new_non_inlined(1)) + ); + assert_eq!( + group_values.get_indices_by_hash(5).unwrap(), + (vec![5, 6, 7], GroupIndexView::new_non_inlined(2)) + ); + assert_eq!(group_values.map.len(), 4); + + // Emit first 1 to test case 4, and cases 5~6 again + let _ = group_values.emit(EmitTo::First(1)).unwrap(); + assert_eq!( + group_values.get_indices_by_hash(2).unwrap(), + (vec![0], GroupIndexView::new_inlined(0)) + ); + assert_eq!( + group_values.get_indices_by_hash(3).unwrap(), + (vec![1], GroupIndexView::new_inlined(1)) + ); + assert_eq!( + group_values.get_indices_by_hash(4).unwrap(), + (vec![2, 3], GroupIndexView::new_non_inlined(0)) + ); + assert_eq!( + group_values.get_indices_by_hash(5).unwrap(), + (vec![4, 5, 6], GroupIndexView::new_non_inlined(1)) + ); + assert_eq!(group_values.map.len(), 4); + + // Emit first 5 to test cases 1~3 again + let _ = group_values.emit(EmitTo::First(5)).unwrap(); + assert_eq!( + group_values.get_indices_by_hash(5).unwrap(), + (vec![0, 1], GroupIndexView::new_non_inlined(0)) + ); + assert_eq!(group_values.map.len(), 1); + + // Emit first 1 to test cases 4 again + let _ = group_values.emit(EmitTo::First(1)).unwrap(); + assert_eq!( + group_values.get_indices_by_hash(5).unwrap(), + (vec![0], GroupIndexView::new_inlined(0)) + ); + assert_eq!(group_values.map.len(), 1); + + // Emit first 1 to test cases 7 + let _ = group_values.emit(EmitTo::First(1)).unwrap(); + assert!(group_values.map.is_empty()); + } + + /// Test data set for [`GroupValuesColumn::vectorized_intern`] + /// + /// Define the test data and support loading them into test [`GroupValuesColumn::vectorized_intern`] + /// + /// The covering situations: + /// + /// Array type: + /// - Primitive array + /// - String(byte) array + /// - String view(byte view) array + /// + /// Repeation and nullability in single batch: + /// - All not null rows + /// - Mixed null + not null rows + /// - All null rows + /// - All not null rows(repeated) + /// - Null + not null rows(repeated) + /// - All not null rows(repeated) + /// + /// If group exists in `map`: + /// - Group exists in inlined group view + /// - Group exists in non-inlined group view + /// - Group not exist + bucket not found in `map` + /// - Group not exist + not equal to inlined group view(tested in hash collision) + /// - Group not exist + not equal to non-inlined group view(tested in hash collision) + struct VectorizedTestDataSet { + test_batches: Vec<Vec<ArrayRef>>, + expected_batch: RecordBatch, + } + + impl VectorizedTestDataSet { + fn new() -> Self { + // Intern batch 1 + let col1 = Int64Array::from(vec![ + // Repeated rows in batch + Some(42), // all not nulls + repeated rows + exist in map case + None, // mixed + repeated rows + exist in map case + None, // mixed + repeated rows + not exist in map case + Some(1142), // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some(42), + None, + None, + Some(1142), + None, + // Unique rows in batch + Some(4211), // all not nulls + unique rows + exist in map case + None, // mixed + unique rows + exist in map case + None, // mixed + unique rows + not exist in map case + Some(4212), // mixed + unique rows + not exist in map case + ]); + + let col2 = StringArray::from(vec![ + // Repeated rows in batch + Some("string1"), // all not nulls + repeated rows + exist in map case + None, // mixed + repeated rows + exist in map case + Some("string2"), // mixed + repeated rows + not exist in map case + None, // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some("string1"), + None, + Some("string2"), + None, + None, + // Unique rows in batch + Some("string3"), // all not nulls + unique rows + exist in map case + None, // mixed + unique rows + exist in map case + Some("string4"), // mixed + unique rows + not exist in map case + None, // mixed + unique rows + not exist in map case + ]); + + let col3 = StringViewArray::from(vec![ + // Repeated rows in batch + Some("stringview1"), // all not nulls + repeated rows + exist in map case + Some("stringview2"), // mixed + repeated rows + exist in map case + None, // mixed + repeated rows + not exist in map case + None, // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some("stringview1"), + Some("stringview2"), + None, + None, + None, + // Unique rows in batch + Some("stringview3"), // all not nulls + unique rows + exist in map case + Some("stringview4"), // mixed + unique rows + exist in map case + None, // mixed + unique rows + not exist in map case + None, // mixed + unique rows + not exist in map case + ]); + let batch1 = vec![ + Arc::new(col1) as _, + Arc::new(col2) as _, + Arc::new(col3) as _, + ]; + + // Intern batch 2 + let col1 = Int64Array::from(vec![ + // Repeated rows in batch + Some(42), // all not nulls + repeated rows + exist in map case + None, // mixed + repeated rows + exist in map case + None, // mixed + repeated rows + not exist in map case + Some(21142), // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some(42), + None, + None, + Some(21142), + None, + // Unique rows in batch + Some(4211), // all not nulls + unique rows + exist in map case + None, // mixed + unique rows + exist in map case + None, // mixed + unique rows + not exist in map case + Some(24212), // mixed + unique rows + not exist in map case + ]); + + let col2 = StringArray::from(vec![ + // Repeated rows in batch + Some("string1"), // all not nulls + repeated rows + exist in map case + None, // mixed + repeated rows + exist in map case + Some("2string2"), // mixed + repeated rows + not exist in map case + None, // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some("string1"), + None, + Some("2string2"), + None, + None, + // Unique rows in batch + Some("string3"), // all not nulls + unique rows + exist in map case + None, // mixed + unique rows + exist in map case + Some("2string4"), // mixed + unique rows + not exist in map case + None, // mixed + unique rows + not exist in map case + ]); + + let col3 = StringViewArray::from(vec![ + // Repeated rows in batch + Some("stringview1"), // all not nulls + repeated rows + exist in map case + Some("stringview2"), // mixed + repeated rows + exist in map case + None, // mixed + repeated rows + not exist in map case + None, // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some("stringview1"), + Some("stringview2"), + None, + None, + None, + // Unique rows in batch + Some("stringview3"), // all not nulls + unique rows + exist in map case + Some("stringview4"), // mixed + unique rows + exist in map case + None, // mixed + unique rows + not exist in map case + None, // mixed + unique rows + not exist in map case + ]); + let batch2 = vec![ + Arc::new(col1) as _, + Arc::new(col2) as _, + Arc::new(col3) as _, + ]; + + // Intern batch 3 + let col1 = Int64Array::from(vec![ + // Repeated rows in batch + Some(42), // all not nulls + repeated rows + exist in map case + None, // mixed + repeated rows + exist in map case + None, // mixed + repeated rows + not exist in map case + Some(31142), // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some(42), + None, + None, + Some(31142), + None, + // Unique rows in batch + Some(4211), // all not nulls + unique rows + exist in map case + None, // mixed + unique rows + exist in map case + None, // mixed + unique rows + not exist in map case + Some(34212), // mixed + unique rows + not exist in map case + ]); + + let col2 = StringArray::from(vec![ + // Repeated rows in batch + Some("string1"), // all not nulls + repeated rows + exist in map case + None, // mixed + repeated rows + exist in map case + Some("3string2"), // mixed + repeated rows + not exist in map case + None, // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some("string1"), + None, + Some("3string2"), + None, + None, + // Unique rows in batch + Some("string3"), // all not nulls + unique rows + exist in map case + None, // mixed + unique rows + exist in map case + Some("3string4"), // mixed + unique rows + not exist in map case + None, // mixed + unique rows + not exist in map case + ]); + + let col3 = StringViewArray::from(vec![ + // Repeated rows in batch + Some("stringview1"), // all not nulls + repeated rows + exist in map case + Some("stringview2"), // mixed + repeated rows + exist in map case + None, // mixed + repeated rows + not exist in map case + None, // mixed + repeated rows + not exist in map case + None, // all nulls + repeated rows + exist in map case + Some("stringview1"), + Some("stringview2"), + None, + None, + None, + // Unique rows in batch + Some("stringview3"), // all not nulls + unique rows + exist in map case + Some("stringview4"), // mixed + unique rows + exist in map case + None, // mixed + unique rows + not exist in map case + None, // mixed + unique rows + not exist in map case + ]); + let batch3 = vec![ + Arc::new(col1) as _, + Arc::new(col2) as _, + Arc::new(col3) as _, + ]; + + // Expected batch + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Utf8View, true), + ])); + + let col1 = Int64Array::from(vec![ + // Repeated rows in batch + Some(42), + None, + None, + Some(1142), + None, + Some(21142), + None, + Some(31142), + None, + // Unique rows in batch + Some(4211), + None, + None, + Some(4212), + None, + Some(24212), + None, + Some(34212), + ]); + + let col2 = StringArray::from(vec![ + // Repeated rows in batch + Some("string1"), + None, + Some("string2"), + None, + Some("2string2"), + None, + Some("3string2"), + None, + None, + // Unique rows in batch + Some("string3"), + None, + Some("string4"), + None, + Some("2string4"), + None, + Some("3string4"), + None, + ]); + + let col3 = StringViewArray::from(vec![ + // Repeated rows in batch + Some("stringview1"), + Some("stringview2"), + None, + None, + None, + None, + None, + None, + None, + // Unique rows in batch + Some("stringview3"), + Some("stringview4"), + None, + None, + None, + None, + None, + None, + ]); + let expected_batch = vec![ + Arc::new(col1) as _, + Arc::new(col2) as _, + Arc::new(col3) as _, + ]; + let expected_batch = RecordBatch::try_new(schema, expected_batch).unwrap(); + + Self { + test_batches: vec![batch1, batch2, batch3], + expected_batch, + } + } + + fn load_to_group_values(&self, group_values: &mut impl GroupValues) { + for batch in self.test_batches.iter() { + group_values.intern(batch, &mut vec![]).unwrap(); + } + } + + fn schema(&self) -> SchemaRef { + self.expected_batch.schema() + } + } + + fn check_result(actual_batch: &RecordBatch, expected_batch: &RecordBatch) { + let formatted_actual_batch = + pretty_format_batches(std::slice::from_ref(actual_batch)) + .unwrap() + .to_string(); + let mut formatted_actual_batch_sorted: Vec<&str> = + formatted_actual_batch.trim().lines().collect(); + formatted_actual_batch_sorted.sort_unstable(); + + let formatted_expected_batch = + pretty_format_batches(std::slice::from_ref(expected_batch)) + .unwrap() + .to_string(); + + let mut formatted_expected_batch_sorted: Vec<&str> = + formatted_expected_batch.trim().lines().collect(); + formatted_expected_batch_sorted.sort_unstable(); + + for (i, (actual_line, expected_line)) in formatted_actual_batch_sorted + .iter() + .zip(&formatted_expected_batch_sorted) + .enumerate() + { + assert_eq!( + (i, actual_line), + (i, expected_line), + "Inconsistent result\n\n\ + Actual batch:\n{formatted_actual_batch}\n\ + Expected batch:\n{formatted_expected_batch}\n\ + ", + ); + } + } + + fn insert_inline_group_index_view( + group_values: &mut GroupValuesColumn<false>, + hash_key: u64, + group_index: u64, + ) { + let group_index_view = GroupIndexView::new_inlined(group_index); + group_values.map.insert_accounted( + (hash_key, group_index_view), + |(hash, _)| *hash, + &mut group_values.map_size, + ); + } + + fn insert_non_inline_group_index_view( + group_values: &mut GroupValuesColumn<false>, + hash_key: u64, + group_indices: Vec<usize>, + ) { + let list_offset = group_values.group_index_lists.len(); + let group_index_view = GroupIndexView::new_non_inlined(list_offset as u64); + group_values.group_index_lists.push(group_indices); + group_values.map.insert_accounted( + (hash_key, group_index_view), + |(hash, _)| *hash, + &mut group_values.map_size, + ); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs new file mode 100644 index 0000000..148c569 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
@@ -0,0 +1,654 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::HashValue; +use crate::aggregates::group_values::multi_group_by::{ + GroupColumn, Nulls, nulls_equal_to, +}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::ArrowNativeTypeOp; +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, BooleanBufferBuilder, PrimitiveArray, + cast::AsArray, +}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::DataType; +use arrow::util::bit_util::apply_bitwise_binary_op; +use datafusion_common::Result; +use datafusion_common::utils::split_vec_min_alloc; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use std::iter; +use std::sync::Arc; + +/// An implementation of [`GroupColumn`] for primitive values +/// +/// Optimized to skip null buffer construction if the input is known to be non nullable +/// +/// # Template parameters +/// +/// `T`: the native Rust type that stores the data +/// `NULLABLE`: if the data can contain any nulls +#[derive(Debug)] +pub struct PrimitiveGroupValueBuilder<T: ArrowPrimitiveType, const NULLABLE: bool> { + data_type: DataType, + group_values: Vec<T::Native>, + nulls: MaybeNullBufferBuilder, +} + +impl<T, const NULLABLE: bool> PrimitiveGroupValueBuilder<T, NULLABLE> +where + T: ArrowPrimitiveType, + T::Native: HashValue, +{ + /// Create a new `PrimitiveGroupValueBuilder` + pub fn new(data_type: DataType) -> Self { + Self { + data_type, + group_values: vec![], + nulls: MaybeNullBufferBuilder::new(), + } + } + + fn vectorized_equal_to_non_nullable( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + assert!( + !NULLABLE || (array.null_count() == 0 && !self.nulls.might_have_nulls()), + "called with nullable input" + ); + let array_values = array.as_primitive::<T>().values(); + let n = lhs_rows.len(); + + // Build a packed comparison bitmask, then AND it into equal_to_results + let num_bytes = n.div_ceil(8); + let mut cmp_buf = vec![0u8; num_bytes]; + + for (i, (&lhs_row, &rhs_row)) in lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(i) { + continue; + } + let left = if cfg!(debug_assertions) { + self.group_values[lhs_row] + } else { + unsafe { *self.group_values.get_unchecked(lhs_row) } + }; + let right = if cfg!(debug_assertions) { + array_values[rhs_row] + } else { + unsafe { *array_values.get_unchecked(rhs_row) } + }; + // `left` was already canonicalized on append; canonicalize the + // input so ±0 (and any future equivalence class) compares equal. + if left.is_eq(right.canonicalize()) { + cmp_buf[i / 8] |= 1 << (i % 8); + } + } + + // AND the comparison result into the existing equal_to_results bitmask + apply_bitwise_binary_op( + equal_to_results.as_slice_mut(), + 0, + &cmp_buf, + 0, + n, + |a, b| a & b, + ); + } + + pub fn vectorized_equal_nullable( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + assert!(NULLABLE, "called with non-nullable input"); + let array = array.as_primitive::<T>(); + + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + if !equal_to_results.get_bit(idx) { + continue; + } + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + if !result { + equal_to_results.set_bit(idx, false); + } + continue; + } + + if !self.group_values[lhs_row].is_eq(array.value(rhs_row).canonicalize()) { + equal_to_results.set_bit(idx, false); + } + } + } +} + +impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn + for PrimitiveGroupValueBuilder<T, NULLABLE> +where + T::Native: HashValue, +{ + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + // Perf: skip null check (by short circuit) if input is not nullable + if NULLABLE { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + // Otherwise, we need to check their values + } + + self.group_values[lhs_row] + .is_eq(array.as_primitive::<T>().value(rhs_row).canonicalize()) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + // Perf: skip null check if input can't have nulls + if NULLABLE { + if array.is_null(row) { + self.nulls.append(true); + self.group_values.push(T::default_value()); + } else { + self.nulls.append(false); + self.group_values + .push(array.as_primitive::<T>().value(row).canonicalize()); + } + } else { + self.group_values + .push(array.as_primitive::<T>().value(row).canonicalize()); + } + + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + if !NULLABLE || (array.null_count() == 0 && !self.nulls.might_have_nulls()) { + self.vectorized_equal_to_non_nullable( + lhs_rows, + array, + rhs_rows, + equal_to_results, + ); + } else { + self.vectorized_equal_nullable(lhs_rows, array, rhs_rows, equal_to_results); + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let arr = array.as_primitive::<T>(); + + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match (NULLABLE, all_null_or_non_null) { + (true, Nulls::Some) => { + for &row in rows { + if array.is_null(row) { + self.nulls.append(true); + self.group_values.push(T::default_value()); + } else { + self.nulls.append(false); + self.group_values.push(arr.value(row).canonicalize()); + } + } + } + + (true, Nulls::None) => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.group_values.push(arr.value(row).canonicalize()); + } + } + + (true, Nulls::All) => { + self.nulls.append_n(rows.len(), true); + self.group_values + .extend(iter::repeat_n(T::default_value(), rows.len())); + } + + (false, _) => { + for &row in rows { + self.group_values.push(arr.value(row).canonicalize()); + } + } + } + + Ok(()) + } + + fn len(&self) -> usize { + self.group_values.len() + } + + fn size(&self) -> usize { + self.group_values.allocated_size() + self.nulls.allocated_size() + } + + fn build(self: Box<Self>) -> ArrayRef { + let Self { + data_type, + group_values, + nulls, + } = *self; + + let nulls = nulls.build(); + if !NULLABLE { + assert!(nulls.is_none(), "unexpected nulls in non nullable input"); + } + + let arr = PrimitiveArray::<T>::new(ScalarBuffer::from(group_values), nulls); + // Set timezone information for timestamp + Arc::new(arr.with_data_type(data_type)) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + let first_n = split_vec_min_alloc(&mut self.group_values, n); + let first_n_nulls = if NULLABLE { self.nulls.take_n(n) } else { None }; + + Arc::new( + PrimitiveArray::<T>::new(ScalarBuffer::from(first_n), first_n_nulls) + .with_data_type(self.data_type.clone()), + ) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; + use arrow::array::{ + ArrayRef, BooleanBufferBuilder, Float32Array, Int32Array, Int64Array, + NullBufferBuilder, + }; + use arrow::datatypes::{DataType, Float32Type, Int32Type, Int64Type}; + + use super::GroupColumn; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec<bool> { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + #[test] + fn test_nullable_primitive_equal_to() { + let append = |builder: &mut PrimitiveGroupValueBuilder<Float32Type, true>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &PrimitiveGroupValueBuilder<Float32Type, true>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_nullable_primitive_equal_to_internal(append, equal_to); + } + + #[test] + fn test_nullable_primitive_vectorized_equal_to() { + let append = |builder: &mut PrimitiveGroupValueBuilder<Float32Type, true>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &PrimitiveGroupValueBuilder<Float32Type, true>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_nullable_primitive_equal_to_internal(append, equal_to); + } + + fn test_nullable_primitive_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut PrimitiveGroupValueBuilder<Float32Type, true>, &ArrayRef, &[usize]), + E: FnMut( + &PrimitiveGroupValueBuilder<Float32Type, true>, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; values not equal + // - exist not null, input not null; values equal + + // Define PrimitiveGroupValueBuilder + let mut builder = + PrimitiveGroupValueBuilder::<Float32Type, true>::new(DataType::Float32); + let builder_array = Arc::new(Float32Array::from(vec![ + None, + None, + None, + Some(1.0), + Some(2.0), + Some(f32::NAN), + Some(3.0), + ])) as ArrayRef; + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5, 6]); + + // Define input array + let (_, values, _nulls) = Float32Array::from(vec![ + Some(1.0), + Some(2.0), + None, + Some(1.0), + None, + Some(f32::NAN), + None, + ]) + .into_parts(); + + // explicitly build a null buffer where one of the null values also happens to match + let mut nulls = NullBufferBuilder::new(6); + nulls.append_non_null(); + nulls.append_null(); + nulls.append_null(); + nulls.append_non_null(); + nulls.append_null(); + nulls.append_non_null(); + nulls.append_null(); + let input_array = Arc::new(Float32Array::new(values, nulls.finish())) as ArrayRef; + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5, 6], + &input_array, + &[0, 1, 2, 3, 4, 5, 6], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(!results[4]); + assert!(results[5]); + assert!(!results[6]); + } + + #[test] + fn test_not_nullable_primitive_equal_to() { + let append = |builder: &mut PrimitiveGroupValueBuilder<Int64Type, false>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &PrimitiveGroupValueBuilder<Int64Type, false>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_not_nullable_primitive_equal_to_internal(append, equal_to); + } + + #[test] + fn test_not_nullable_primitive_vectorized_equal_to() { + let append = |builder: &mut PrimitiveGroupValueBuilder<Int64Type, false>, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &PrimitiveGroupValueBuilder<Int64Type, false>, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_not_nullable_primitive_equal_to_internal(append, equal_to); + } + + fn test_not_nullable_primitive_equal_to_internal<A, E>(mut append: A, mut equal_to: E) + where + A: FnMut(&mut PrimitiveGroupValueBuilder<Int64Type, false>, &ArrayRef, &[usize]), + E: FnMut( + &PrimitiveGroupValueBuilder<Int64Type, false>, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - values equal + // - values not equal + + // Define PrimitiveGroupValueBuilder + let mut builder = + PrimitiveGroupValueBuilder::<Int64Type, false>::new(DataType::Int64); + let builder_array = + Arc::new(Int64Array::from(vec![Some(0), Some(1)])) as ArrayRef; + append(&mut builder, &builder_array, &[0, 1]); + + // Define input array + let input_array = Arc::new(Int64Array::from(vec![Some(0), Some(2)])) as ArrayRef; + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1], + &input_array, + &[0, 1], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(!results[1]); + } + + #[test] + fn test_nullable_primitive_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = + PrimitiveGroupValueBuilder::<Int64Type, true>::new(DataType::Int64); + + // All nulls input array + let all_nulls_input_array = Arc::new(Int64Array::from(vec![ + Option::<i64>::None, + None, + None, + None, + None, + ])) as _; + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = Arc::new(Int64Array::from(vec![ + Some(1), + Some(2), + Some(3), + Some(4), + Some(5), + ])) as _; + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } + + // All bits false: every row must be skipped; accessing any lhs/rhs index would panic. + #[test] + fn test_vectorized_equal_to_skips_false_rows() { + let mut builder = + PrimitiveGroupValueBuilder::<Int32Type, true>::new(DataType::Int32); + let array = Arc::new(Int32Array::from(vec![None::<i32>, None])) as ArrayRef; + builder.vectorized_append(&array, &[0, 1]).unwrap(); + + let mut results = BooleanBufferBuilder::new(2); + results.append_n(2, false); + + builder.vectorized_equal_to( + &[usize::MAX, usize::MAX], + &array, + &[usize::MAX, usize::MAX], + &mut results, + ); + } + + #[test] + fn test_primitive_take_n() { + // drain branch: n * 2 <= len + let mut builder = + PrimitiveGroupValueBuilder::<Int64Type, true>::new(DataType::Int64); + let array = Arc::new(Int64Array::from(vec![ + Some(10), + None, + Some(30), + Some(40), + Some(50), + ])) as ArrayRef; + for i in 0..5 { + builder.append_val(&array, i).unwrap(); + } + // len=5, n=2, n*2=4 <= 5 → drain branch + let out = builder.take_n(2); + let expected = Arc::new(Int64Array::from(vec![Some(10), None])) as ArrayRef; + assert_eq!(&out, &expected); + // remaining: [30, 40, 50] + assert_eq!(builder.len(), 3); + + // split_off branch: remaining < n (len=3, n=2, n*2=4 > 3) + let out2 = builder.take_n(2); + let expected2 = Arc::new(Int64Array::from(vec![Some(30), Some(40)])) as ArrayRef; + assert_eq!(&out2, &expected2); + // remaining: [50] + assert_eq!(builder.len(), 1); + + // take the last element + let out3 = builder.take_n(1); + let expected3 = Arc::new(Int64Array::from(vec![Some(50)])) as ArrayRef; + assert_eq!(&out3, &expected3); + assert_eq!(builder.len(), 0); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs new file mode 100644 index 0000000..1445a81 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs
@@ -0,0 +1,1129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A generic [`GroupColumn`] backed by the arrow row format. +//! +//! Unlike the type-specialized builders in this module (primitive, byte, +//! boolean, ...), [`RowsGroupColumn`] works for *any* data type that arrow's +//! [`RowConverter`] can encode — including nested types such as `Struct`, +//! `List`, `LargeList` and `FixedSizeList`. It stores one group value per row +//! in a single-column [`Rows`] buffer and compares group keys by their encoded +//! bytes. +//! +//! # Why this exists +//! +//! [`GroupValuesColumn`] can only be used when *every* column of the group-by +//! key has a [`GroupColumn`] implementation; otherwise the whole aggregation +//! falls back to the row-wise [`GroupValuesRows`], which is materially slower +//! and heavier for the columns that *would* have qualified for the column-wise +//! fast path. By providing a generic fallback `GroupColumn`, a schema like +//! `GROUP BY int_col, struct_col` keeps `int_col` on its fast native builder +//! and only pays the row-encoding cost on `struct_col`, instead of dragging both +//! columns onto `GroupValuesRows`. +//! +//! # Relationship to hashing +//! +//! This column does not hash anything itself: [`GroupValuesColumn`] hashes the +//! raw input columns via `create_hashes`, which already supports nested types. +//! Equality is decided here by comparing arrow-row bytes. For the two to agree +//! on group identity, values that this column considers equal must hash equal — +//! see the float `-0.0` / `NaN` note on [`RowsGroupColumn`]. +//! +//! [`GroupValuesColumn`]: crate::aggregates::group_values::multi_group_by::GroupValuesColumn +//! [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows + +use crate::aggregates::group_values::multi_group_by::GroupColumn; +use crate::aggregates::group_values::row::encode_array_if_necessary; + +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::datatypes::DataType; +use arrow::row::{RowConverter, Rows, SortField}; +use datafusion_common::{DataFusionError, Result}; + +/// A [`GroupColumn`] that stores group values for a single column in the arrow +/// [row format], backed by a single-field [`RowConverter`]. +/// +/// # NULL semantics +/// +/// The [`GroupColumn`] contract treats two NULLs as equal. The row format +/// encodes NULL with a distinct sentinel, so `null`-row bytes compare equal to +/// each other and unequal to any non-null row — matching the contract without +/// special-casing. +/// +/// # Float `-0.0` / `NaN` +/// +/// Equality here is byte equality under arrow's IEEE-754 *totalOrder* row +/// encoding, which treats `-0.0` and `+0.0` as distinct and canonicalizes +/// `NaN`. Because hashing is performed separately (on the raw input array), a +/// caller must ensure the two agree — e.g. by normalizing `-0.0 → +0.0` on the +/// input columns before hashing when a float leaf is present (as +/// [`GroupValuesRows`] does). See the module docs. +/// +/// [row format]: arrow::row +/// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows +pub struct RowsGroupColumn { + /// Single-field row converter for this column's data type. + row_converter: RowConverter, + /// Accumulated group values in row format; `group_values.row(i)` is the + /// group value for group index `i`. + group_values: Rows, + /// The column's expected output type. The row format decodes dictionary / + /// run-end encoded values to their plain value type, so emitted arrays are + /// re-encoded to this type in `build` / `take_n` (mirroring + /// `GroupValuesRows::emit`). + output_type: DataType, +} + +/// Walk `data_type`'s subtree and return `true` if it contains a +/// [`DataType::FixedSizeList`] whose descendant tree includes any +/// [`DataType::Dictionary`]. +/// +/// Two-state recursion: once we cross a `FixedSizeList`, `inside_fsl` +/// stays true for every descendant, so a `Dictionary` anywhere below +/// counts. Above that boundary, encountering a `Dictionary` is fine — +/// only nested containers propagate the risk. +/// +/// TODO: this guard works around <https://github.com/apache/arrow-rs/issues/10413> +/// (`decode_fixed_size_list` panics instead of applying the +/// dictionary-flatten `corrected_type` step). Fixed upstream by +/// <https://github.com/apache/arrow-rs/pull/10414> (merged 2026-07-24, not +/// yet in a release as of arrow 59.1.0). Once DataFusion upgrades to an +/// arrow release containing that fix, `FixedSizeList<Dictionary>` will +/// decode like the other list-likes (flattened child, re-encoded by +/// `encode_array_if_necessary`'s existing `FixedSizeList` arm) — remove +/// this guard and its `supports_type` rejection at that point. +fn contains_fsl_with_dictionary(data_type: &DataType) -> bool { + fn walk(dt: &DataType, inside_fsl: bool) -> bool { + match dt { + DataType::Dictionary(_, _) => inside_fsl, + DataType::FixedSizeList(f, _) => walk(f.data_type(), true), + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) => walk(f.data_type(), inside_fsl), + DataType::Map(f, _) => walk(f.data_type(), inside_fsl), + DataType::Struct(fs) => fs.iter().any(|f| walk(f.data_type(), inside_fsl)), + DataType::RunEndEncoded(_, values) => walk(values.data_type(), inside_fsl), + DataType::Union(fs, _) => { + fs.iter().any(|(_, f)| walk(f.data_type(), inside_fsl)) + } + _ => false, + } + } + walk(data_type, false) +} + +/// Return `true` if `data_type` contains a [`DataType::Union`] or +/// [`DataType::RunEndEncoded`] anywhere in its subtree. +/// +/// These two nested variants can round-trip through `RowConverter` in +/// principle, but their arrow-row decoders have not been validated by +/// this crate's test matrix against the full range of leaf types (dict, +/// nested, etc.). Before this PR both were handled by `GroupValuesRows` +/// (they were not `is_nested`-eligible for `GroupValuesColumn`), so +/// reject them here to preserve the pre-PR routing rather than route +/// untested shapes through `RowsGroupColumn`. When we grow explicit +/// round-trip tests for these types, this blacklist can be removed. +fn contains_union_or_run_end_encoded(data_type: &DataType) -> bool { + match data_type { + DataType::Union(_, _) | DataType::RunEndEncoded(_, _) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) => { + contains_union_or_run_end_encoded(f.data_type()) + } + DataType::Map(f, _) => contains_union_or_run_end_encoded(f.data_type()), + DataType::Struct(fs) => fs + .iter() + .any(|f| contains_union_or_run_end_encoded(f.data_type())), + _ => false, + } +} + +impl RowsGroupColumn { + /// Returns whether `data_type` can be handled by this generic column. + /// + /// This is stricter than [`RowConverter::supports_fields`]: the row + /// format also has to survive the `build` / `take_n` reverse trip + /// through [`RowConverter::convert_rows`], and arrow's + /// `decode_fixed_size_list` (arrow-row 59.1.0) skips the + /// dictionary-flatten correction that the other list-like decoders + /// apply, so any `FixedSizeList` containing a `Dictionary` leaf + /// panics on emit with `"FixedSizeListArray expected data type + /// Dictionary(...) got <flattened> for \"item\""`. + /// + /// Reject those shapes here so `make_group_column` falls back to + /// `GroupValuesRows`. The other list-likes (`List`, `LargeList`, + /// `ListView`, `LargeListView`, `Map`) do carry the correction, so + /// they decode without panicking — but the correction *flattens* any + /// dictionary child to its value type, so `build` / `take_n` must + /// re-encode the emitted array back to `output_type` via + /// `encode_array_if_necessary` (which has a reconstruction arm for + /// each of these containers). + /// + /// Additionally, `Union` and `RunEndEncoded` are rejected because + /// they were routed to `GroupValuesRows` before this column existed + /// and their arrow-row round-trip has not been covered by this + /// crate's tests yet. Keeping them on the pre-PR path avoids + /// introducing an untested code path for those types. + pub fn supports_type(data_type: &DataType) -> bool { + if contains_fsl_with_dictionary(data_type) { + return false; + } + if contains_union_or_run_end_encoded(data_type) { + return false; + } + RowConverter::supports_fields(&[SortField::new(data_type.clone())]) + } + + /// Create an empty [`RowsGroupColumn`] for `data_type`. + pub fn try_new(data_type: DataType) -> Result<Self> { + let row_converter = RowConverter::new(vec![SortField::new(data_type.clone())])?; + let group_values = row_converter.empty_rows(0, 0); + Ok(Self { + row_converter, + group_values, + output_type: data_type, + }) + } + + /// Materialize `rows` into a single array of `self.output_type`, re-applying + /// dictionary / run-end encoding the row format strips on decode. + fn rows_to_array<'a>( + &self, + rows: impl IntoIterator<Item = arrow::row::Row<'a>>, + ) -> ArrayRef { + let mut arrays = self + .row_converter + .convert_rows(rows) + .expect("row conversion during emit"); + assert_eq!( + arrays.len(), + 1, + "Single field row converter must produce exactly one array, actual length is {}", + arrays.len() + ); + let array = arrays.pop().unwrap(); + encode_array_if_necessary(&array, &self.output_type) + .expect("dictionary re-encode during emit") + } + + /// Encode a whole incoming column into the row format. + fn convert(&self, array: &ArrayRef) -> Result<Rows> { + self.row_converter + .convert_columns(std::slice::from_ref(array)) + .map_err(DataFusionError::from) + } +} + +impl GroupColumn for RowsGroupColumn { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + // Scalar path (hash-collision remainder / streaming). Encode just the + // single incoming row rather than the whole column. The vectorized + // methods below encode the batch once; this path is expected to be rare. + let incoming = self + .convert(&array.slice(rhs_row, 1)) + .expect("row conversion during equal_to"); + self.group_values.row(lhs_row) == incoming.row(0) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let incoming = self.convert(&array.slice(row, 1))?; + self.group_values.push(incoming.row(0)); + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + // Encode the incoming column once for the whole batch. + let incoming = self + .convert(array) + .expect("row conversion during vectorized_equal_to"); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + // Preserve the AND-accumulate contract: skip rows already false. + if !equal_to_results.get_bit(idx) { + continue; + } + if self.group_values.row(lhs_row) != incoming.row(rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + // Encode the incoming column once, then push the selected rows. + let incoming = self.convert(array)?; + for &row in rows { + self.group_values.push(incoming.row(row)); + } + Ok(()) + } + + fn len(&self) -> usize { + self.group_values.num_rows() + } + + fn size(&self) -> usize { + self.row_converter.size() + self.group_values.size() + } + + fn build(self: Box<Self>) -> ArrayRef { + self.rows_to_array(&self.group_values) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(n <= self.group_values.num_rows()); + + // Materialize the first `n` group rows. + let output = self.rows_to_array(self.group_values.iter().take(n)); + + // Shift the remaining rows to the front by rebuilding the buffer. + // TODO: mirror the arrow-rs efficiency TODO in `GroupValuesRows::emit`. + let remaining_rows = self.group_values.num_rows() - n; + let remaining_bytes = self.group_values.lengths().skip(n).sum(); + let mut remaining = self + .row_converter + .empty_rows(remaining_rows, remaining_bytes); + for row in self.group_values.iter().skip(n) { + remaining.push(row); + } + self.group_values = remaining; + + output + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{ + Array, ArrayRef, FixedSizeListArray, Int32Array, StringArray, StructArray, + }; + use arrow::datatypes::{DataType, Field, Int32Type}; + use std::sync::Arc; + + fn fsl_i32(data: Vec<Option<Vec<Option<i32>>>>, list_len: i32) -> ArrayRef { + Arc::new(FixedSizeListArray::from_iter_primitive::<Int32Type, _, _>( + data, list_len, + )) + } + + /// Build a `FixedSizeList<Utf8>` with `list_len == 1`. Each entry is one + /// row holding a single (optionally null) string, and an outer `None` + /// marks a null list. Variable-length string payloads give retained rows + /// distinct encoded lengths, which is what `take_n`'s byte preallocation + /// depends on. + fn fsl_utf8(rows: Vec<Option<Option<&str>>>) -> ArrayRef { + let child = StringArray::from( + rows.iter() + .map(|row| row.and_then(|inner| inner)) + .collect::<Vec<_>>(), + ); + let outer_nulls = arrow::buffer::NullBuffer::from( + rows.iter().map(|row| row.is_some()).collect::<Vec<_>>(), + ); + Arc::new(FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Utf8, true)), + 1, + Arc::new(child), + Some(outer_nulls), + )) + } + + /// The generic column must agree with a per-row reference for equality, + /// including inner-null and outer-null rows, on a `FixedSizeList<Int32>`. + #[test] + fn fsl_append_equal_to_build_roundtrip() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 2, + ); + let mut col = Box::new(RowsGroupColumn::try_new(dt).unwrap()); + + // group values: [1,2], null-outer, [3, null-inner] + let input = fsl_i32( + vec![ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3), None]), + ], + 2, + ); + + col.vectorized_append(&input, &[0, 1, 2]).unwrap(); + assert_eq!(col.len(), 3); + + // Probe with a fresh batch: row0 == group0, row1 (null) == group1, + // row2 differs from group0, row3 (inner null) == group2. + let probe = fsl_i32( + vec![ + Some(vec![Some(1), Some(2)]), // == g0 + None, // == g1 + Some(vec![Some(9), Some(9)]), // != g0 + Some(vec![Some(3), None]), // == g2 + ], + 2, + ); + + assert!(col.equal_to(0, &probe, 0)); + assert!(col.equal_to(1, &probe, 1)); + assert!(!col.equal_to(0, &probe, 2)); + assert!(col.equal_to(2, &probe, 3)); + + // Vectorized equal_to should match the scalar reference. + let mut results = BooleanBufferBuilder::new(3); + results.append_n(3, true); + col.vectorized_equal_to(&[0, 1, 2], &probe, &[0, 1, 3], &mut results); + assert!(results.get_bit(0)); + assert!(results.get_bit(1)); + assert!(results.get_bit(2)); + + // build() must reproduce the original group values. + let out = col.build(); + let out = out.as_any().downcast_ref::<FixedSizeListArray>().unwrap(); + assert_eq!(out.len(), 3); + assert!(out.is_null(1)); + assert!(!out.is_null(0)); + } + + /// `take_n` must emit the first `n` rows and shift the rest to the front. + #[test] + fn fsl_take_n_shifts_remaining() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + let input = fsl_i32( + vec![ + Some(vec![Some(10)]), + Some(vec![Some(20)]), + Some(vec![Some(30)]), + ], + 1, + ); + col.vectorized_append(&input, &[0, 1, 2]).unwrap(); + + let first = col.take_n(1); + let first = first.as_any().downcast_ref::<FixedSizeListArray>().unwrap(); + let first_vals = first + .value(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap() + .clone(); + assert_eq!(first_vals.value(0), 10); + assert_eq!(col.len(), 2); + + // Remaining 20, 30 should now be at indices 0, 1. + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::<FixedSizeListArray>().unwrap(); + assert_eq!(rest.len(), 2); + let g0 = rest + .value(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap() + .value(0); + assert_eq!(g0, 20); + } + + /// `take_n` preallocates the retained-row buffer from the known retained + /// row count and byte size + /// + /// To exercise the byte-sum path directly, the retained rows are + /// `FixedSizeList<Utf8>` values with deliberately unequal payload + /// lengths plus an inner-null. Here we assert every emitted and + /// every shifted-down value is byte-for-byte unchanged. + #[test] + fn take_n_preallocated_rebuild_preserves_variable_length_rows() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Utf8, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + // Rows 0-2 are emitted; rows 3-6 are retained and shifted to the + // front. The retained rows intentionally have different encoded + // lengths so `lengths().skip(3).sum()` is not a simple row_count * k. + let input = fsl_utf8(vec![ + Some(Some("emit_a")), // 0: emitted + Some(None), // 1: emitted (inner-null) + None, // 2: emitted (outer-null) + Some(Some("")), // 3: retained, empty payload + Some(Some("xyz")), // 4: retained, short payload + Some(None), // 5: retained, inner-null + Some(Some("a_much_longer_payload_string")), // 6: retained, long payload + ]); + col.vectorized_append(&input, &[0, 1, 2, 3, 4, 5, 6]) + .unwrap(); + assert_eq!(col.len(), 7); + + // Emit the first three rows; four rows should remain. + let emitted = col.take_n(3); + let emitted = emitted + .as_any() + .downcast_ref::<FixedSizeListArray>() + .unwrap(); + assert_eq!(emitted.len(), 3); + assert_eq!( + emitted + .value(0) + .as_any() + .downcast_ref::<StringArray>() + .unwrap() + .value(0), + "emit_a" + ); + // Row 1 was an inner-null; row 2 was an outer-null. + assert!( + emitted + .value(1) + .as_any() + .downcast_ref::<StringArray>() + .unwrap() + .is_null(0) + ); + assert!(emitted.is_null(2)); + + assert_eq!(col.len(), 4); + + // The four retained rows must survive the rebuild intact, in order: + // "", "xyz", inner-null, "a_much_longer_payload_string". + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::<FixedSizeListArray>().unwrap(); + assert_eq!(rest.len(), 4); + + let value_at = |idx: usize| { + rest.value(idx) + .as_any() + .downcast_ref::<StringArray>() + .unwrap() + .clone() + }; + assert_eq!(value_at(0).value(0), ""); + assert_eq!(value_at(1).value(0), "xyz"); + assert!( + value_at(2).is_null(0), + "retained inner-null row must be preserved" + ); + assert_eq!(value_at(3).value(0), "a_much_longer_payload_string"); + } + + /// Works for `Struct<a: Int32>` too — proves the column is type-generic. + #[test] + fn struct_roundtrip() { + let dt = DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2)])); + let input: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("a", DataType::Int32, true)].into(), + vec![a], + None, + )); + col.vectorized_append(&input, &[0, 1]).unwrap(); + assert_eq!(col.len(), 2); + assert!(col.equal_to(0, &input, 0)); + assert!(!col.equal_to(0, &input, 1)); + } + + #[test] + fn supports_type_matches_row_converter_impl() { + assert!(RowsGroupColumn::supports_type(&DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 3 + ))); + assert!(RowsGroupColumn::supports_type(&DataType::Struct( + vec![Field::new("a", DataType::Int32, true)].into() + ))); + // Whether Map is encodable depends on the arrow-rs version. + // Just assert that our `supports_type` agrees with arrow's + // `RowConverter::supports_fields` — either both accept it or both + // reject it. Both are correct wrt the invariant. + let map_field = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int32, true), + ] + .into(), + ), + false, + )); + let map_dt = DataType::Map(map_field, false); + let arrow_supports = + RowConverter::supports_fields(&[SortField::new(map_dt.clone())]); + assert_eq!(RowsGroupColumn::supports_type(&map_dt), arrow_supports); + } + + /// Regression test for the nested-container recursion in + /// [`crate::aggregates::group_values::row::encode_array_if_necessary`]. + /// `RowConverter` flattens dictionary values on the way in, so a + /// `List<Dict<Int32, Utf8>>` schema round-trips with `Utf8` values + /// unless the helper re-encodes the leaf. Without that recursion, + /// `build()` would emit an array whose data type does not match the + /// group column's declared type. + #[test] + fn build_preserves_list_of_dictionary_schema() { + use arrow::array::{DictionaryArray, ListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Int32Type; + + let dict_dt = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let item_field = Arc::new(Field::new("item", dict_dt.clone(), true)); + let outer_dt = DataType::List(Arc::clone(&item_field)); + + // Skip if this arrow-rs version rejects the nesting — the invariant we + // care about is `output().data_type() == declared type` conditional on + // supports_type saying yes. + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + // Build List<Dict<Int32,Utf8>> of one row = ["a", "b"]. + let values = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1]); + let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap(); + let offsets = OffsetBuffer::from_lengths([2]); + let list = + ListArray::try_new(Arc::clone(&item_field), offsets, Arc::new(dict), None) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + col.vectorized_append(&input, &[0]).unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "build() must return the declared List<Dict> data type, \ + not the RowConverter-flattened List<Utf8>", + ); + } + + // ---- FSL<Dict> rejection ---------------------------------------- + // + // arrow-row 59.1.0's `decode_fixed_size_list` skips the + // dict-flatten correction that the generic `decode<L>` path applies + // to `List` / `LargeList` / `ListView` / `LargeListView` / `Map`, + // so any `FixedSizeList` containing a `Dictionary` leaf panics on + // emit. `supports_type` must reject those shapes so + // `GroupValuesRows` fallback handles them instead. These tests pin + // the current shape of that black-list. + + fn dict_utf8() -> DataType { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + + fn fsl_of(inner: DataType) -> DataType { + DataType::FixedSizeList(Arc::new(Field::new("item", inner, true)), 2) + } + + #[test] + fn supports_type_rejects_fixed_size_list_of_dict() { + // Direct case: `FixedSizeList<Dict<Int32, Utf8>>`. + assert!(!RowsGroupColumn::supports_type(&fsl_of(dict_utf8()))); + } + + #[test] + fn supports_type_rejects_fsl_with_dict_nested_in_struct() { + // The dict is one level deep under a struct that is itself the + // FSL element. arrow-row still panics because `convert_raw` + // returns the struct with a decoded (Utf8) field while the + // FSL builder expects the declared struct-with-dict shape. + let struct_dt = DataType::Struct(vec![Field::new("d", dict_utf8(), true)].into()); + assert!(!RowsGroupColumn::supports_type(&fsl_of(struct_dt))); + } + + #[test] + fn supports_type_rejects_fsl_with_dict_nested_in_list() { + // `FixedSizeList<List<Dict>>` — the inner `List<Dict>` handles + // dicts correctly on its own, but the outer FSL wrapper still + // panics with the mismatched declared child type. + let list_of_dict = + DataType::List(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(!RowsGroupColumn::supports_type(&fsl_of(list_of_dict))); + } + + #[test] + fn supports_type_rejects_fsl_hidden_under_outer_list() { + // Sibling positioning: the outer container is a `List` (which is + // fine on its own), but its child is a `FixedSizeList<Dict>`. + // The panic surface is at the inner FSL layer regardless of what + // wraps it, so this must still be rejected. + let outer = + DataType::List(Arc::new(Field::new("item", fsl_of(dict_utf8()), true))); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + #[test] + fn supports_type_rejects_fsl_hidden_under_outer_struct() { + // Same, but the outer wrapper is a struct. + let outer = + DataType::Struct(vec![Field::new("f", fsl_of(dict_utf8()), true)].into()); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + // ---- FSL without dicts is still fine ---------------------------- + + #[test] + fn supports_type_accepts_fsl_of_primitive() { + // Sanity: a plain FSL<Int32> must not get caught by the + // dict-under-FSL blacklist. + assert!(RowsGroupColumn::supports_type(&fsl_of(DataType::Int32))); + } + + #[test] + fn supports_type_accepts_fsl_of_struct_without_dict() { + // FSL of struct where the struct's fields are all primitives. + let struct_dt = + DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + assert!(RowsGroupColumn::supports_type(&fsl_of(struct_dt))); + } + + // ---- Positive round-trip tests for non-FSL list-likes ----------- + // + // The other list-like decoders in arrow-row 59.1.0 + // (`GenericListArrayOrMap` path) apply the corrected_type fix, so + // `List<Dict>`, `LargeList<Dict>`, `ListView<Dict>`, `LargeListView<Dict>` + // and `Map<..., Dict>` all round-trip cleanly. These tests pin + // that they are (a) accepted by `supports_type` and (b) actually + // survive `vectorized_append` + `build()` without panicking, so a + // future arrow-rs regression there is caught here rather than in + // production. + + #[test] + fn supports_type_accepts_large_list_of_dict() { + let dt = DataType::LargeList(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_accepts_list_view_of_dict() { + let dt = DataType::ListView(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_accepts_large_list_view_of_dict() { + let dt = DataType::LargeListView(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_map_agrees_with_row_converter() { + // Map<Int32, Dict<Int32, Utf8>>. Whether arrow-row supports Map + // depends on the version; either way, our `supports_type` must + // agree with `RowConverter::supports_fields` — otherwise we'd + // pick a strategy the converter can't back. + let entries = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", dict_utf8(), true), + ] + .into(), + ), + false, + )); + let map_dt = DataType::Map(entries, false); + let arrow_supports = + RowConverter::supports_fields(&[SortField::new(map_dt.clone())]); + assert_eq!(RowsGroupColumn::supports_type(&map_dt), arrow_supports); + } + + /// End-to-end regression: `LargeList<Dict<Int32, Utf8>>` must + /// actually survive `vectorized_append` + `build()` on the current + /// arrow-rs version, not just be accepted by `supports_type`. + #[test] + fn build_preserves_large_list_of_dictionary_schema() { + use arrow::array::{DictionaryArray, LargeListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::LargeList(Arc::clone(&item_field)); + + // Skip if this arrow-rs version rejects the nesting (defensive: + // the invariant we care about is `output().data_type() == declared` + // conditional on `supports_type` saying yes). + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + let values = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1]); + let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap(); + let offsets = OffsetBuffer::<i64>::from_lengths([2]); + let list = LargeListArray::try_new( + Arc::clone(&item_field), + offsets, + Arc::new(dict), + None, + ) + .unwrap(); + + col.vectorized_append(&(Arc::new(list) as ArrayRef), &[0]) + .unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "LargeList<Dict>: build() must preserve the declared type", + ); + } + + /// Build a two-row `ListView<Dict<Int32, Utf8>>` array with rows + /// `["a", "b"]` and `["c"]` — the shape from the review reproducer: + /// `arrow_cast(a, 'ListView(Dictionary(Int32, Utf8))')`. + fn list_view_of_dict_input() -> (DataType, ArrayRef) { + use arrow::array::{DictionaryArray, ListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::ListView(Arc::clone(&item_field)); + + let values = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2]); + let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::<i32>::from(vec![0, 2]); + let sizes = ScalarBuffer::<i32>::from(vec![2, 1]); + let list = ListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + (outer_dt, Arc::new(list) as ArrayRef) + } + + /// `ListView<Dict>`: arrow-row's `decode_list_view` flattens the + /// dictionary child (`corrected_type`), so `build` must re-encode + /// the emitted array back to the declared type. Regression for the + /// review reproducer that failed with + /// `expected ListView(Dictionary(Int32, Utf8)) but found ListView(Utf8)`. + #[test] + fn build_preserves_list_view_of_dictionary_schema() { + let (outer_dt, input) = list_view_of_dict_input(); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + assert_eq!(col.len(), 2); + + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "ListView<Dict>: build() must return the declared type, \ + not the RowConverter-flattened ListView<Utf8>", + ); + assert_eq!(built.len(), 2); + } + + /// Same regression through the `take_n` path (used by + /// `EmitTo::First(n)`), including the type of the *remaining* + /// values emitted by a subsequent `build`. + #[test] + fn take_n_preserves_list_view_of_dictionary_schema() { + let (outer_dt, input) = list_view_of_dict_input(); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let taken = col.take_n(1); + assert_eq!( + taken.data_type(), + &outer_dt, + "ListView<Dict>: take_n() must return the declared type", + ); + assert_eq!(taken.len(), 1); + + let rest = col.build(); + assert_eq!( + rest.data_type(), + &outer_dt, + "ListView<Dict>: build() after take_n must also preserve the type", + ); + assert_eq!(rest.len(), 1); + } + + /// `LargeListView<Dict>` fails the same way as `ListView<Dict>` + /// per the review; cover both `build` and `take_n`. + #[test] + fn build_and_take_n_preserve_large_list_view_of_dictionary_schema() { + use arrow::array::{DictionaryArray, LargeListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::LargeListView(Arc::clone(&item_field)); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let values = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2]); + let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::<i64>::from(vec![0, 2]); + let sizes = ScalarBuffer::<i64>::from(vec![2, 1]); + let list = LargeListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let taken = col.take_n(1); + assert_eq!( + taken.data_type(), + &outer_dt, + "LargeListView<Dict>: take_n() must return the declared type", + ); + + let rest = col.build(); + assert_eq!( + rest.data_type(), + &outer_dt, + "LargeListView<Dict>: build() must return the declared type", + ); + assert_eq!(rest.len(), 1); + } + + /// Group-identity must survive the dictionary flatten + re-encode + /// round trip: appending the same logical list twice (with distinct + /// dictionary key mappings) must map to one group, a different list + /// to another. Mirrors the review reproducer's GROUP BY semantics + /// (2 distinct groups from 3 input rows). + #[test] + fn list_view_of_dict_groups_by_logical_value() { + use arrow::array::{DictionaryArray, ListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::ListView(Arc::clone(&item_field)); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + // Rows: ["a","b"], ["a","b"], ["c"] → 2 distinct groups. + let values = Arc::new(StringArray::from(vec!["a", "b", "a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2, 3, 4]); + let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::<i32>::from(vec![0, 2, 4]); + let sizes = ScalarBuffer::<i32>::from(vec![2, 2, 1]); + let list = ListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + // Append row 0 as group 0. + col.vectorized_append(&input, &[0]).unwrap(); + // Row 1 must compare equal to group 0 (same logical value). + assert!( + col.equal_to(0, &input, 1), + "identical logical lists must be equal regardless of dict keys", + ); + // Row 2 must not. + assert!( + !col.equal_to(0, &input, 2), + "different logical lists must not be equal", + ); + + col.vectorized_append(&input, &[2]).unwrap(); + assert_eq!(col.len(), 2, "3 input rows → 2 distinct groups"); + + let built = col.build(); + assert_eq!(built.data_type(), &outer_dt); + assert_eq!(built.len(), 2); + } + + /// End-to-end regression for `Map<Int32, Dict<Int32, Utf8>>` when + /// arrow-row supports it. Same intent as the LargeList test. + #[test] + fn build_preserves_map_of_dictionary_schema() { + use arrow::array::{ + DictionaryArray, Int32Array, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + + let key_field = Arc::new(Field::new("keys", DataType::Int32, false)); + let value_field = Arc::new(Field::new("values", dict_utf8(), true)); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(vec![(*key_field).clone(), (*value_field).clone()].into()), + false, + )); + let outer_dt = DataType::Map(Arc::clone(&entries_field), false); + + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + // One map entry: {1 -> "a"}. + let keys = Arc::new(Int32Array::from(vec![1])) as ArrayRef; + let values_arr = Arc::new(StringArray::from(vec!["a"])); + let value_keys = Int32Array::from(vec![0]); + let value_dict = + DictionaryArray::<Int32Type>::try_new(value_keys, values_arr).unwrap(); + let entries = StructArray::try_new( + vec![(*key_field).clone(), (*value_field).clone()].into(), + vec![keys, Arc::new(value_dict)], + None, + ) + .unwrap(); + let offsets = OffsetBuffer::<i32>::from_lengths([1]); + let map = + MapArray::try_new(Arc::clone(&entries_field), offsets, entries, None, false) + .unwrap(); + + col.vectorized_append(&(Arc::new(map) as ArrayRef), &[0]) + .unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "Map<..., Dict>: build() must preserve the declared type", + ); + } + + // ---- Union / RunEndEncoded defensive rejection ----------------- + // + // Before this PR both types were routed to `GroupValuesRows` + // (`group_column_supported_type` didn't have a nested branch). This + // PR added `is_nested`-based dispatch to `RowsGroupColumn`, which + // would opt them in — but the arrow-row round-trip for these two + // families hasn't been covered by our tests. Reject them here so + // the pre-PR routing is preserved; drop the blacklist when the + // round-trip matrix grows to include them. + + #[test] + fn supports_type_rejects_union() { + use arrow::datatypes::UnionFields; + + let fields = UnionFields::try_new( + vec![0_i8, 1_i8], + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ], + ) + .unwrap(); + let dt = DataType::Union(fields, arrow::datatypes::UnionMode::Dense); + assert!( + !RowsGroupColumn::supports_type(&dt), + "Union must fall back to GroupValuesRows until arrow-row \ + round-trip is covered by our tests", + ); + } + + #[test] + fn supports_type_rejects_run_end_encoded_with_nested_values() { + // REE with `is_nested() = true` (nested values) is what this PR + // could otherwise opt into RowsGroupColumn; keep it on + // GroupValuesRows. + let list_of_i32 = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let dt = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", list_of_i32, true)), + ); + assert!(!RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_rejects_run_end_encoded_with_scalar_values() { + // REE with scalar values is `is_nested() == false`, so + // `group_column_supported_type` never routes it to us via the + // nested branch anyway — but pin the invariant explicitly so a + // future refactor doesn't accidentally opt it in. + let dt = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ); + assert!(!RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_rejects_ree_hidden_under_outer_wrapper() { + // REE buried under a struct or list: still rejected because + // the wrapper's decoder recurses through the REE branch we + // haven't validated. + let ree = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ); + let outer = DataType::Struct(vec![Field::new("f", ree, true)].into()); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + #[test] + fn supports_type_accepts_plain_list_and_struct_still() { + // Sanity: the defensive Union/REE blacklist must not accidentally + // catch the well-tested list-likes / structs that this column + // exists to serve. + let list_of_int = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + assert!(RowsGroupColumn::supports_type(&list_of_int)); + + let struct_of_prims = DataType::Struct( + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ] + .into(), + ); + assert!(RowsGroupColumn::supports_type(&struct_of_prims)); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/null_builder.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/null_builder.rs new file mode 100644 index 0000000..6a84d68 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/null_builder.rs
@@ -0,0 +1,100 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::NullBufferBuilder; +use arrow::buffer::NullBuffer; + +/// Builder for an (optional) null mask +/// +/// Optimized for avoid creating the bitmask when all values are non-null +#[derive(Debug)] +pub(crate) struct MaybeNullBufferBuilder { + /// Note this is an Arrow *VALIDITY* buffer (so it is false for nulls, true + /// for non-nulls) + nulls: NullBufferBuilder, +} + +impl MaybeNullBufferBuilder { + /// Create a new builder + pub fn new() -> Self { + Self { + nulls: NullBufferBuilder::new(0), + } + } + + /// Return true if the row at index `row` is null + pub fn is_null(&self, row: usize) -> bool { + match self.nulls.as_slice() { + // validity mask means a unset bit is NULL + Some(_) => !self.nulls.is_valid(row), + None => false, + } + } + + /// Set the nullness of the next row to `is_null` + /// + /// If `value` is true, the row is null. + /// If `value` is false, the row is non null + pub fn append(&mut self, is_null: bool) { + self.nulls.append(!is_null) + } + + pub fn append_n(&mut self, n: usize, is_null: bool) { + if is_null { + self.nulls.append_n_nulls(n); + } else { + self.nulls.append_n_non_nulls(n); + } + } + + /// return the number of heap allocated bytes used by this structure to store boolean values + pub fn allocated_size(&self) -> usize { + // NullBufferBuilder builder::allocated_size returns capacity in bits + self.nulls.allocated_size() / 8 + } + + /// Return a NullBuffer representing the accumulated nulls so far + pub fn build(mut self) -> Option<NullBuffer> { + self.nulls.finish() + } + + /// Returns a NullBuffer representing the first `n` rows accumulated so far + /// shifting any remaining down by `n` + pub fn take_n(&mut self, n: usize) -> Option<NullBuffer> { + // Copy over the values at n..len-1 values to the start of a + // new builder and leave it in self + // + // TODO: it would be great to use something like `set_bits` from arrow here. + let mut new_builder = NullBufferBuilder::new(self.nulls.len()); + for i in n..self.nulls.len() { + new_builder.append(self.nulls.is_valid(i)); + } + std::mem::swap(&mut new_builder, &mut self.nulls); + + // take only first n values from the original builder + new_builder.truncate(n); + new_builder.finish() + } + + /// Returns true if this builder might have any nulls + /// + /// This is guaranteed to be true if there are nulls + /// but may be true even if there are no nulls + pub(crate) fn might_have_nulls(&self) -> bool { + self.nulls.as_slice().is_some() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs new file mode 100644 index 0000000..cbd7a60 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/row.rs
@@ -0,0 +1,414 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::GroupValues; +use arrow::array::{ + Array, ArrayRef, FixedSizeListArray, LargeListArray, LargeListViewArray, ListArray, + ListViewArray, MapArray, PrimitiveArray, RunArray, StructArray, + downcast_run_end_index, +}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, SchemaRef}; +use arrow::row::{RowConverter, Rows, SortField}; +use datafusion_common::Result; +use datafusion_common::hash_utils::RandomState; +use datafusion_common::hash_utils::create_hashes; +use datafusion_common::utils::normalize_float_zero; +use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; +use datafusion_expr::EmitTo; +use hashbrown::hash_table::HashTable; +use log::debug; +use std::mem::size_of; +use std::sync::Arc; + +/// A [`GroupValues`] making use of [`Rows`] +/// +/// This is a general implementation of [`GroupValues`] that works for any +/// combination of data types and number of columns, including nested types such as +/// structs and lists. +/// +/// It uses the arrow-rs [`Rows`] to store the group values, which is a row-wise +/// representation. +pub struct GroupValuesRows { + /// The output schema + schema: SchemaRef, + + /// Converter for the group values + row_converter: RowConverter, + + /// Logically maps group values to a group_index in + /// [`Self::group_values`] and in each accumulator + /// + /// Uses the raw API of hashbrown to avoid actually storing the + /// keys (group values) in the table + /// + /// keys: u64 hashes of the GroupValue + /// values: (hash, group_index) + map: HashTable<(u64, usize)>, + + /// The size of `map` in bytes + map_size: usize, + + /// The actual group by values, stored in arrow [`Row`] format. + /// `group_values[i]` holds the group value for group_index `i`. + /// + /// The row format is used to compare group keys quickly and store + /// them efficiently in memory. Quick comparison is especially + /// important for multi-column group keys. + /// + /// [`Row`]: arrow::row::Row + group_values: Option<Rows>, + + /// reused buffer to store hashes + hashes_buffer: Vec<u64>, + + /// reused buffer to store rows + rows_buffer: Rows, + + /// Random state for creating hashes + random_state: RandomState, +} + +impl GroupValuesRows { + pub fn try_new(schema: SchemaRef) -> Result<Self> { + // Print a debugging message, so it is clear when the (slower) fallback + // GroupValuesRows is used. + debug!("Creating GroupValuesRows for schema: {schema}"); + let row_converter = RowConverter::new( + schema + .fields() + .iter() + .map(|f| SortField::new(f.data_type().clone())) + .collect(), + )?; + + let map = HashTable::with_capacity(0); + + let starting_rows_capacity = 1000; + + let starting_data_capacity = 64 * starting_rows_capacity; + let rows_buffer = + row_converter.empty_rows(starting_rows_capacity, starting_data_capacity); + Ok(Self { + schema, + row_converter, + map, + map_size: 0, + group_values: None, + hashes_buffer: Default::default(), + rows_buffer, + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + }) + } +} + +impl GroupValues for GroupValuesRows { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and + // primitive hashing both group ±0 together. No-op for non-float + // columns. + let normalized_cols: Vec<ArrayRef> = + cols.iter().map(normalize_float_zero).collect(); + let cols = normalized_cols.as_slice(); + + // Convert the group keys into the row format + let group_rows = &mut self.rows_buffer; + group_rows.clear(); + self.row_converter.append(group_rows, cols)?; + let n_rows = group_rows.num_rows(); + + let mut group_values = match self.group_values.take() { + Some(group_values) => group_values, + None => self.row_converter.empty_rows(0, 0), + }; + + // tracks to which group each of the input rows belongs + groups.clear(); + + // 1.1 Calculate the group keys for the group values + let batch_hashes = &mut self.hashes_buffer; + batch_hashes.clear(); + batch_hashes.resize(n_rows, 0); + create_hashes(cols, &self.random_state, batch_hashes)?; + + for (row, &target_hash) in batch_hashes.iter().enumerate() { + let entry = self.map.find_mut(target_hash, |(exist_hash, group_idx)| { + // Somewhat surprisingly, this closure can be called even if the + // hash doesn't match, so check the hash first with an integer + // comparison first avoid the more expensive comparison with + // group value. https://github.com/apache/datafusion/pull/11718 + target_hash == *exist_hash + // verify that the group that we are inserting with hash is + // actually the same key value as the group in + // existing_idx (aka group_values @ row) + && group_rows.row(row) == group_values.row(*group_idx) + }); + + let group_idx = match entry { + // Existing group_index for this group value + Some((_hash, group_idx)) => *group_idx, + // 1.2 Need to create new entry for the group + None => { + // Add new entry to aggr_state and save newly created index + let group_idx = group_values.num_rows(); + group_values.push(group_rows.row(row)); + + // for hasher function, use precomputed hash value + self.map.insert_accounted( + (target_hash, group_idx), + |(hash, _group_index)| *hash, + &mut self.map_size, + ); + group_idx + } + }; + groups.push(group_idx); + } + + self.group_values = Some(group_values); + + Ok(()) + } + + fn size(&self) -> usize { + let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0); + self.row_converter.size() + + group_values_size + + self.map_size + + self.rows_buffer.size() + + self.hashes_buffer.allocated_size() + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { + self.group_values + .as_ref() + .map(|group_values| group_values.num_rows()) + .unwrap_or(0) + } + + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + let mut group_values = self + .group_values + .take() + .expect("Can not emit from empty rows"); + + let mut output = match emit_to { + EmitTo::All => { + let output = self.row_converter.convert_rows(&group_values)?; + group_values.clear(); + self.map.clear(); + output + } + EmitTo::First(n) => { + let groups_rows = group_values.iter().take(n); + let output = self.row_converter.convert_rows(groups_rows)?; + // Clear out first n group keys by copying them to a new Rows. + // TODO file some ticket in arrow-rs to make this more efficient? + let mut new_group_values = self.row_converter.empty_rows(0, 0); + for row in group_values.iter().skip(n) { + new_group_values.push(row); + } + std::mem::swap(&mut new_group_values, &mut group_values); + + self.map.retain(|(_exists_hash, group_idx)| { + // Decrement group index by n + match group_idx.checked_sub(n) { + // Group index was >= n, shift value down + Some(sub) => { + *group_idx = sub; + true + } + // Group index was < n, so remove from table + None => false, + } + }); + output + } + }; + + // TODO: Materialize dictionaries in group keys + // https://github.com/apache/datafusion/issues/7647 + for (field, array) in self.schema.fields.iter().zip(&mut output) { + let expected = field.data_type(); + *array = encode_array_if_necessary(array, expected)?; + } + + self.group_values = Some(group_values); + Ok(output) + } + + fn clear_shrink(&mut self, num_rows: usize) { + self.group_values = self.group_values.take().map(|mut rows| { + rows.clear(); + rows + }); + self.map.clear(); + self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); + self.hashes_buffer.clear(); + self.hashes_buffer.shrink_to(num_rows); + } +} + +/// Re-apply dictionary / run-end encoding to `array` so it matches `expected`. +/// +/// Arrow's [`RowConverter`] flattens dictionary and run-end-encoded values to +/// their plain value type during row encoding (at [`RowConverter::append`]), +/// so any group-value array produced from the row format is in that plain +/// type and must be re-encoded to match the schema's expected type before +/// being returned. Shared with the generic row-backed `GroupColumn`. +/// +/// [`RowConverter`]: arrow::row::RowConverter +/// [`RowConverter::append`]: arrow::row::RowConverter::append +pub(crate) fn encode_array_if_necessary( + array: &ArrayRef, + expected: &DataType, +) -> Result<ArrayRef> { + match (expected, array.data_type()) { + (DataType::Struct(expected_fields), _) => { + let struct_array = array.as_any().downcast_ref::<StructArray>().unwrap(); + let arrays = expected_fields + .iter() + .zip(struct_array.columns()) + .map(|(expected_field, column)| { + encode_array_if_necessary(column, expected_field.data_type()) + }) + .collect::<Result<Vec<_>>>()?; + + Ok(Arc::new(StructArray::try_new( + expected_fields.clone(), + arrays, + struct_array.nulls().cloned(), + )?)) + } + (DataType::List(expected_field), &DataType::List(_)) => { + let list = array.as_any().downcast_ref::<ListArray>().unwrap(); + + Ok(Arc::new(ListArray::try_new( + Arc::<arrow::datatypes::Field>::clone(expected_field), + list.offsets().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::LargeList(expected_field), &DataType::LargeList(_)) => { + let list = array.as_any().downcast_ref::<LargeListArray>().unwrap(); + + Ok(Arc::new(LargeListArray::try_new( + Arc::<arrow::datatypes::Field>::clone(expected_field), + list.offsets().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::ListView(expected_field), &DataType::ListView(_)) => { + // arrow-row's `decode_list_view` applies the dictionary-flatten + // `corrected_type` to the child, so a `ListView<Dictionary<..>>` + // decodes as `ListView<value type>` and the child must be + // re-encoded here (same as `List` above, plus the `sizes` + // buffer that view-lists carry). + let list = array.as_any().downcast_ref::<ListViewArray>().unwrap(); + + Ok(Arc::new(ListViewArray::try_new( + Arc::<arrow::datatypes::Field>::clone(expected_field), + list.offsets().clone(), + list.sizes().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::LargeListView(expected_field), &DataType::LargeListView(_)) => { + let list = array.as_any().downcast_ref::<LargeListViewArray>().unwrap(); + + Ok(Arc::new(LargeListViewArray::try_new( + Arc::<arrow::datatypes::Field>::clone(expected_field), + list.offsets().clone(), + list.sizes().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + ( + DataType::FixedSizeList(expected_field, expected_size), + &DataType::FixedSizeList(_, _), + ) => { + let list = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap(); + + Ok(Arc::new(FixedSizeListArray::try_new( + Arc::<arrow::datatypes::Field>::clone(expected_field), + *expected_size, + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::Map(expected_entries_field, ordered), &DataType::Map(_, _)) => { + let map = array.as_any().downcast_ref::<MapArray>().unwrap(); + // Re-encode the entries `StructArray` (which holds key/value + // columns) against the expected entries field's struct type. + let entries_as_ref: ArrayRef = Arc::new(map.entries().clone()); + let entries = encode_array_if_necessary( + &entries_as_ref, + expected_entries_field.data_type(), + )?; + let entries = entries + .as_any() + .downcast_ref::<StructArray>() + .expect("Map entries recurse must yield a StructArray") + .clone(); + Ok(Arc::new(MapArray::try_new( + Arc::<arrow::datatypes::Field>::clone(expected_entries_field), + map.offsets().clone(), + entries, + map.nulls().cloned(), + *ordered, + )?)) + } + (DataType::Dictionary(_, _), _) => Ok(cast(array.as_ref(), expected)?), + ( + DataType::RunEndEncoded(run_ends_field, expected_values_field), + &DataType::RunEndEncoded(_, _), + ) => { + macro_rules! reencode_ree { + ($run_end_type:ty) => {{ + let run_array = array + .as_any() + .downcast_ref::<RunArray<$run_end_type>>() + .unwrap(); + let values = encode_array_if_necessary( + &(Arc::clone(run_array.values()) as ArrayRef), + expected_values_field.data_type(), + )?; + let run_ends = PrimitiveArray::<$run_end_type>::new( + run_array.run_ends().inner().clone(), + None, + ); + Ok(Arc::new(RunArray::try_new(&run_ends, &values)?)) + }}; + } + downcast_run_end_index! { + run_ends_field.data_type() => (reencode_ree), + _ => unreachable!("unsupported run end type: {}", run_ends_field.data_type()), + } + } + (DataType::RunEndEncoded(_, _), _) => Ok(cast(array.as_ref(), expected)?), + (_, _) => Ok(Arc::<dyn Array>::clone(array)), + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/boolean.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/boolean.rs new file mode 100644 index 0000000..e993c0c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/boolean.rs
@@ -0,0 +1,153 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::GroupValues; + +use arrow::array::{ + ArrayRef, AsArray as _, BooleanArray, BooleanBufferBuilder, NullBufferBuilder, +}; +use datafusion_common::Result; +use datafusion_expr::EmitTo; +use std::{mem::size_of, sync::Arc}; + +#[derive(Debug)] +pub struct GroupValuesBoolean { + false_group: Option<usize>, + true_group: Option<usize>, + null_group: Option<usize>, +} + +impl GroupValuesBoolean { + pub fn new() -> Self { + Self { + false_group: None, + true_group: None, + null_group: None, + } + } +} + +impl GroupValues for GroupValuesBoolean { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + let array = cols[0].as_boolean(); + groups.clear(); + + for value in array.iter() { + let index = match value { + Some(false) => { + if let Some(index) = self.false_group { + index + } else { + let index = self.len(); + self.false_group = Some(index); + index + } + } + Some(true) => { + if let Some(index) = self.true_group { + index + } else { + let index = self.len(); + self.true_group = Some(index); + index + } + } + None => { + if let Some(index) = self.null_group { + index + } else { + let index = self.len(); + self.null_group = Some(index); + index + } + } + }; + + groups.push(index); + } + + Ok(()) + } + + fn size(&self) -> usize { + size_of::<Self>() + } + + fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn len(&self) -> usize { + self.false_group.is_some() as usize + + self.true_group.is_some() as usize + + self.null_group.is_some() as usize + } + + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + let len = self.len(); + let mut builder = BooleanBufferBuilder::new(len); + let emit_count = match emit_to { + EmitTo::All => len, + EmitTo::First(n) => n, + }; + builder.append_n(emit_count, false); + if let Some(idx) = self.true_group.as_mut() { + if *idx < emit_count { + builder.set_bit(*idx, true); + self.true_group = None; + } else { + *idx -= emit_count; + } + } + + if let Some(idx) = self.false_group.as_mut() { + if *idx < emit_count { + // already false, no need to set + self.false_group = None; + } else { + *idx -= emit_count; + } + } + + let values = builder.finish(); + + let nulls = if let Some(idx) = self.null_group.as_mut() { + if *idx < emit_count { + let mut buffer = NullBufferBuilder::new(len); + buffer.append_n_non_nulls(*idx); + buffer.append_null(); + buffer.append_n_non_nulls(emit_count - *idx - 1); + + self.null_group = None; + Some(buffer.finish().unwrap()) + } else { + *idx -= emit_count; + None + } + } else { + None + }; + + Ok(vec![Arc::new(BooleanArray::new(values, nulls)) as _]) + } + + fn clear_shrink(&mut self, _num_rows: usize) { + self.false_group = None; + self.true_group = None; + self.null_group = None; + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs new file mode 100644 index 0000000..b881a51 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes.rs
@@ -0,0 +1,128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::mem::size_of; + +use crate::aggregates::group_values::GroupValues; + +use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; +use datafusion_common::Result; +use datafusion_expr::EmitTo; +use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; + +/// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values +/// +/// This specialization is significantly faster than using the more general +/// purpose `Row`s format +pub struct GroupValuesBytes<O: OffsetSizeTrait> { + /// Map string/binary values to group index + map: ArrowBytesMap<O, usize>, + /// The total number of groups so far (used to assign group_index) + num_groups: usize, +} + +impl<O: OffsetSizeTrait> GroupValuesBytes<O> { + pub fn new(output_type: OutputType) -> Self { + Self { + map: ArrowBytesMap::new(output_type), + num_groups: 0, + } + } +} + +impl<O: OffsetSizeTrait> GroupValues for GroupValuesBytes<O> { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + assert_eq!(cols.len(), 1); + + // look up / add entries in the table + let arr = &cols[0]; + + groups.clear(); + self.map.insert_if_new( + arr, + // called for each new group + |_value| { + // assign new group index on each insert + let group_idx = self.num_groups; + self.num_groups += 1; + group_idx + }, + // called for each group + |group_idx| { + groups.push(group_idx); + }, + ); + + // ensure we assigned a group to for each row + assert_eq!(groups.len(), arr.len()); + Ok(()) + } + + fn size(&self) -> usize { + self.map.size() + size_of::<Self>() + } + + fn is_empty(&self) -> bool { + self.num_groups == 0 + } + + fn len(&self) -> usize { + self.num_groups + } + + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + // Reset the map to default, and convert it into a single array + let map_contents = self.map.take().into_state(); + + let group_values = match emit_to { + EmitTo::All => { + self.num_groups -= map_contents.len(); + map_contents + } + EmitTo::First(n) if n == self.len() => { + self.num_groups -= map_contents.len(); + map_contents + } + EmitTo::First(n) => { + // if we only wanted to take the first n, insert the rest back + // into the map we could potentially avoid this reallocation, at + // the expense of much more complex code. + // see https://github.com/apache/datafusion/issues/9195 + let emit_group_values = map_contents.slice(0, n); + let remaining_group_values = + map_contents.slice(n, map_contents.len() - n); + + self.num_groups = 0; + let mut group_indexes = vec![]; + self.intern(&[remaining_group_values], &mut group_indexes)?; + + // Verify that the group indexes were assigned in the correct order + assert_eq!(0, group_indexes[0]); + + emit_group_values + } + }; + + Ok(vec![group_values]) + } + + 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(); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs new file mode 100644 index 0000000..7a56f7c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs
@@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::GroupValues; +use arrow::array::{Array, ArrayRef}; +use datafusion_expr::EmitTo; +use datafusion_physical_expr::binary_map::OutputType; +use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; +use std::mem::size_of; + +/// A [`GroupValues`] storing single column of Utf8View/BinaryView values +/// +/// This specialization is significantly faster than using the more general +/// purpose `Row`s format +pub struct GroupValuesBytesView { + /// Map string/binary values to group index + map: ArrowBytesViewMap<usize>, + /// The total number of groups so far (used to assign group_index) + num_groups: usize, +} + +impl GroupValuesBytesView { + pub fn new(output_type: OutputType) -> Self { + Self { + map: ArrowBytesViewMap::new(output_type), + num_groups: 0, + } + } +} + +impl GroupValues for GroupValuesBytesView { + fn intern( + &mut self, + cols: &[ArrayRef], + groups: &mut Vec<usize>, + ) -> datafusion_common::Result<()> { + assert_eq!(cols.len(), 1); + + // look up / add entries in the table + let arr = &cols[0]; + + groups.clear(); + self.map.insert_if_new( + arr, + // called for each new group + |_value| { + // assign new group index on each insert + let group_idx = self.num_groups; + self.num_groups += 1; + group_idx + }, + // called for each group + |group_idx| { + groups.push(group_idx); + }, + ); + + // ensure we assigned a group to for each row + assert_eq!(groups.len(), arr.len()); + Ok(()) + } + + fn size(&self) -> usize { + self.map.size() + size_of::<Self>() + } + + fn is_empty(&self) -> bool { + self.num_groups == 0 + } + + fn len(&self) -> usize { + self.num_groups + } + + fn emit(&mut self, emit_to: EmitTo) -> datafusion_common::Result<Vec<ArrayRef>> { + // Reset the map to default, and convert it into a single array + let map_contents = self.map.take().into_state(); + + let group_values = match emit_to { + EmitTo::All => { + self.num_groups -= map_contents.len(); + map_contents + } + EmitTo::First(n) if n == self.len() => { + self.num_groups -= map_contents.len(); + map_contents + } + EmitTo::First(n) => { + // if we only wanted to take the first n, insert the rest back + // into the map we could potentially avoid this reallocation, at + // the expense of much more complex code. + // see https://github.com/apache/datafusion/issues/9195 + let emit_group_values = map_contents.slice(0, n); + let remaining_group_values = + map_contents.slice(n, map_contents.len() - n); + + self.num_groups = 0; + let mut group_indexes = vec![]; + self.intern(&[remaining_group_values], &mut group_indexes)?; + + // Verify that the group indexes were assigned in the correct order + assert_eq!(0, group_indexes[0]); + + emit_group_values + } + }; + + Ok(vec![group_values]) + } + + 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(); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/mod.rs new file mode 100644 index 0000000..89c6b62 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/mod.rs
@@ -0,0 +1,23 @@ +// 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. + +//! `GroupValues` implementations for single group by cases + +pub(crate) mod boolean; +pub(crate) mod bytes; +pub(crate) mod bytes_view; +pub(crate) mod primitive;
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/primitive.rs new file mode 100644 index 0000000..e254aeb --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
@@ -0,0 +1,296 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::GroupValues; +use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; +use arrow::array::{ + ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, NullBufferBuilder, PrimitiveArray, + cast::AsArray, +}; +use arrow::datatypes::{DataType, i256}; +use datafusion_common::Result; +use datafusion_common::hash_utils::RandomState; +use datafusion_common::utils::split_vec_min_alloc; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; +use half::f16; +use hashbrown::hash_table::HashTable; +#[cfg(not(feature = "force_hash_collisions"))] +use std::hash::BuildHasher; +use std::mem::size_of; +use std::sync::Arc; + +/// A trait to allow hashing of floating point numbers +pub trait HashValue { + fn hash(&self, state: &RandomState) -> u64; + + /// Return a canonical representative whose bit pattern is identical for + /// all values that should be grouped together. Default is the identity; + /// floats override this to fold `-0.0` into `+0.0` so the bit-equal + /// `is_eq` check used during insertion treats them as the same group. + /// NaN payload bits are preserved. + #[inline] + fn canonicalize(self) -> Self + where + Self: Sized, + { + self + } +} + +macro_rules! hash_integer { + ($($t:ty),+) => { + $(impl HashValue for $t { + #[cfg(not(feature = "force_hash_collisions"))] + fn hash(&self, state: &RandomState) -> u64 { + state.hash_one(self) + } + + #[cfg(feature = "force_hash_collisions")] + fn hash(&self, _state: &RandomState) -> u64 { + 0 + } + })+ + }; +} +hash_integer!(i8, i16, i32, i64, i128, i256); +hash_integer!(u8, u16, u32, u64); +hash_integer!(IntervalDayTime, IntervalMonthDayNano); + +macro_rules! hash_float { + ($($t:ty),+) => { + $(impl HashValue for $t { + #[cfg(not(feature = "force_hash_collisions"))] + fn hash(&self, state: &RandomState) -> u64 { + state.hash_one(self.canonicalize().to_bits()) + } + + #[cfg(feature = "force_hash_collisions")] + fn hash(&self, _state: &RandomState) -> u64 { + 0 + } + + #[inline] + fn canonicalize(self) -> Self { + let bits = self.to_bits(); + let bits = if bits << 1 == 0 { 0 } else { bits }; + Self::from_bits(bits) + } + })+ + }; +} + +hash_float!(f16, f32, f64); + +/// A [`GroupValues`] storing a single column of primitive values +/// +/// This specialization is significantly faster than using the more general +/// purpose `Row`s format +pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> { + /// The data type of the output array + data_type: DataType, + /// Stores the `(group_index, hash)` based on the hash of its value + /// + /// We also store `hash` is for reducing cost of rehashing. Such cost + /// is obvious in high cardinality group by situation. + /// More details can see: + /// <https://github.com/apache/datafusion/issues/15961> + map: HashTable<(usize, u64)>, + /// The group index of the null value if any + null_group: Option<usize>, + /// The values for each group index + values: Vec<T::Native>, + /// The random state used to generate hashes + random_state: RandomState, +} + +impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> { + pub fn new(data_type: DataType) -> Self { + assert!(PrimitiveArray::<T>::is_compatible(&data_type)); + Self { + data_type, + map: HashTable::with_capacity(128), + values: Vec::with_capacity(128), + null_group: None, + random_state: crate::aggregates::AGGREGATION_HASH_SEED, + } + } +} + +impl<T: ArrowPrimitiveType> GroupValues for GroupValuesPrimitive<T> +where + T::Native: HashValue, +{ + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> { + assert_eq!(cols.len(), 1); + groups.clear(); + + for v in cols[0].as_primitive::<T>() { + let group_id = match v { + None => *self.null_group.get_or_insert_with(|| { + let group_id = self.values.len(); + self.values.push(Default::default()); + group_id + }), + Some(key) => { + // Fold equivalence-class duplicates (e.g. `-0.0` → `+0.0`) + // so the bit-equal `is_eq` matches and the stored value is + // the canonical representative. + let key = key.canonicalize(); + let state = &self.random_state; + let hash = key.hash(state); + let insert = self.map.entry( + hash, + |&(g, h)| unsafe { + hash == h && self.values.get_unchecked(g).is_eq(key) + }, + |&(_, h)| h, + ); + + match insert { + hashbrown::hash_table::Entry::Occupied(o) => o.get().0, + hashbrown::hash_table::Entry::Vacant(v) => { + let g = self.values.len(); + v.insert((g, hash)); + self.values.push(key); + g + } + } + } + }; + groups.push(group_id) + } + Ok(()) + } + + fn size(&self) -> usize { + self.map.capacity() * size_of::<(usize, u64)>() + self.values.allocated_size() + } + + fn is_empty(&self) -> bool { + self.values.is_empty() + } + + fn len(&self) -> usize { + self.values.len() + } + + fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + fn build_primitive<T: ArrowPrimitiveType>( + values: Vec<T::Native>, + null_idx: Option<usize>, + ) -> PrimitiveArray<T> { + let nulls = null_idx.map(|null_idx| { + let mut buffer = NullBufferBuilder::new(values.len()); + buffer.append_n_non_nulls(null_idx); + buffer.append_null(); + buffer.append_n_non_nulls(values.len() - null_idx - 1); + // NOTE: The inner builder must be constructed as there is at least one null + buffer.finish().unwrap() + }); + PrimitiveArray::<T>::new(values.into(), nulls) + } + + let array: PrimitiveArray<T> = match emit_to { + EmitTo::All => { + self.map.clear(); + build_primitive(std::mem::take(&mut self.values), self.null_group.take()) + } + EmitTo::First(n) => { + self.map.retain(|entry| { + // Decrement group index by n + let group_idx = entry.0; + match group_idx.checked_sub(n) { + // Group index was >= n, shift value down + Some(sub) => { + entry.0 = sub; + true + } + // Group index was < n, so remove from table + None => false, + } + }); + let null_group = match &mut self.null_group { + Some(v) if *v >= n => { + *v -= n; + None + } + Some(_) => self.null_group.take(), + None => None, + }; + build_primitive(split_vec_min_alloc(&mut self.values, n), null_group) + } + }; + + Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))]) + } + + fn clear_shrink(&mut self, num_rows: usize) { + self.values.clear(); + self.values.shrink_to(num_rows); + self.map.clear(); + self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::types::Int32Type; + use arrow::array::{ArrayRef, Int32Array}; + use arrow::datatypes::DataType; + use datafusion_expr::EmitTo; + use std::sync::Arc; + + /// Mirror of the `EmitTo::take_needed` regression test, applied to the + /// concrete `GroupValuesPrimitive` accumulator. + /// + /// When `n` is small, the old `split_off(n) + swap` pattern used inside + /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation + /// and returned the emitted prefix carrying the original large backing. + /// + /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: + /// the emitted prefix gets a compact allocation and `self.values` retains the + /// original large one. + #[test] + fn emit_first_small_n_allocates_minimally() -> Result<()> { + let mut gv = GroupValuesPrimitive::<Int32Type>::new(DataType::Int32); + + // Intern 20 distinct values; `new()` pre-allocates capacity 128 for `values`. + let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..20i32)); + let mut groups = vec![]; + gv.intern(&[arr], &mut groups)?; + let capacity_before = gv.values.capacity(); // 128 + + // n=4, n*2=8 <= len=20 -> drain branch + let emitted = gv.emit(EmitTo::First(4))?; + + assert_eq!(emitted[0].len(), 4); + + // `self.values` must retain its original large allocation. + // Old split_off+swap left it with a fresh small allocation (~16). + assert_eq!( + gv.values.capacity(), + capacity_before, + "self.values capacity {} should equal original {} after small First(n) emit", + gv.values.capacity(), + capacity_before, + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/grouped_hash_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/grouped_hash_stream.rs new file mode 100644 index 0000000..99c1011 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/grouped_hash_stream.rs
@@ -0,0 +1,1603 @@ +// 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. + +//! Hash aggregation + +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::vec; + +use super::order::GroupOrdering; +use super::skip_partial::SkipAggregationProbe; +use super::{AggregateExec, format_human_display}; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::order::GroupOrderingFull; +use crate::aggregates::{ + AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, + create_schema, evaluate_group_by, evaluate_many, evaluate_optional, group_id_array, + max_duplicate_ordinal, +}; +use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; +use crate::stream::EmptyRecordBatchStream; +use crate::{PhysicalExpr, aggregates, metrics}; +use crate::{RecordBatchStream, SendableRecordBatchStream}; + +use arrow::array::*; +use arrow::datatypes::SchemaRef; +use datafusion_common::{ + DataFusionError, Result, assert_eq_or_internal_err, assert_or_internal_err, + internal_err, resources_datafusion_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_physical_expr::aggregate::AggregateFunctionExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{GroupsAccumulatorAdapter, PhysicalSortExpr}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; + +use crate::sorts::IncrementalSortIterator; +use datafusion_common::instant::Instant; +use datafusion_common::utils::memory::get_record_batch_memory_size; +use futures::ready; +use futures::stream::{Stream, StreamExt}; +use log::debug; + +#[derive(Debug, Clone)] +/// This object tracks the aggregation phase (input/output) +pub(crate) enum ExecutionState { + ReadingInput, + /// When producing output, the remaining rows to output are stored + /// here and are sliced off as needed in batch_size chunks + ProducingOutput(RecordBatch), + /// Produce intermediate aggregate state for each input row without + /// aggregation. + /// + /// See "partial aggregation" discussion on [`GroupedHashAggregateStream`] + SkippingAggregation, + /// All input has been consumed and all groups have been emitted + Done, +} + +/// This encapsulates the spilling state +struct SpillState { + // ======================================================================== + // PROPERTIES: + // These fields are initialized at the start and remain constant throughout + // the execution. + // ======================================================================== + /// Sorting expression for spilling batches + spill_expr: LexOrdering, + + /// Schema for spilling batches + spill_schema: SchemaRef, + + /// aggregate_arguments for merging spilled data + merging_aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>, + + /// GROUP BY expressions for merging spilled data + merging_group_by: PhysicalGroupBy, + + /// Manages the process of spilling and reading back intermediate data + spill_manager: SpillManager, + + // ======================================================================== + // STATES: + // Fields changes during execution. Can be buffer, or state flags that + // influence the execution in parent `GroupedHashAggregateStream` + // ======================================================================== + /// If data has previously been spilled, the locations of the + /// spill files (in Arrow IPC format) + spills: Vec<SortedSpillFile>, + + /// true when streaming merge is in progress + is_stream_merging: bool, + + // ======================================================================== + // METRICS: + // ======================================================================== + /// Peak memory used for buffered data. + /// Calculated as sum of peak memory values across partitions + peak_mem_used: metrics::Gauge, + // Metrics related to spilling are managed inside `spill_manager` +} + +/// Controls the behavior when an out-of-memory condition occurs. +#[derive(PartialEq, Debug)] +enum OutOfMemoryMode { + /// When out of memory occurs, spill state to disk + Spill, + /// When out of memory occurs, attempt to emit group values early + EmitEarly, + /// When out of memory occurs, immediately report the error + ReportError, +} + +/// HashTable based Grouping Aggregator +/// +/// # Development Note +/// +/// This implementation is being incrementally refactored. See the tracking issue +/// for details. +/// +/// New features and improvements should go directly into the new implementation. +/// Please coordinate through the tracking issue. +/// +/// Issue: <https://github.com/apache/datafusion/issues/22710> +/// +/// # Design Goals +/// +/// This structure is designed so that updating the aggregates can be +/// vectorized (done in a tight loop) without allocations. The +/// accumulator state is *not* managed by this operator (e.g in the +/// hash table) and instead is delegated to the individual +/// accumulators which have type specialized inner loops that perform +/// the aggregation. +/// +/// # Architecture +/// +/// ```text +/// +/// Assigns a consecutive group internally stores aggregate values +/// index for each unique set for all groups +/// of group values +/// +/// ┌────────────┐ ┌──────────────┐ ┌──────────────┐ +/// │ ┌────────┐ │ │┌────────────┐│ │┌────────────┐│ +/// │ │ "A" │ │ ││accumulator ││ ││accumulator ││ +/// │ ├────────┤ │ ││ 0 ││ ││ N ││ +/// │ │ "Z" │ │ ││ ┌────────┐ ││ ││ ┌────────┐ ││ +/// │ └────────┘ │ ││ │ state │ ││ ││ │ state │ ││ +/// │ │ ││ │┌─────┐ │ ││ ... ││ │┌─────┐ │ ││ +/// │ ... │ ││ │├─────┤ │ ││ ││ │├─────┤ │ ││ +/// │ │ ││ │└─────┘ │ ││ ││ │└─────┘ │ ││ +/// │ │ ││ │ │ ││ ││ │ │ ││ +/// │ ┌────────┐ │ ││ │ ... │ ││ ││ │ ... │ ││ +/// │ │ "Q" │ │ ││ │ │ ││ ││ │ │ ││ +/// │ └────────┘ │ ││ │┌─────┐ │ ││ ││ │┌─────┐ │ ││ +/// │ │ ││ │└─────┘ │ ││ ││ │└─────┘ │ ││ +/// └────────────┘ ││ └────────┘ ││ ││ └────────┘ ││ +/// │└────────────┘│ │└────────────┘│ +/// └──────────────┘ └──────────────┘ +/// +/// group_values accumulators +/// +/// ``` +/// +/// For example, given a query like `COUNT(x), SUM(y) ... GROUP BY z`, +/// [`group_values`] will store the distinct values of `z`. There will +/// be one accumulator for `COUNT(x)`, specialized for the data type +/// of `x` and one accumulator for `SUM(y)`, specialized for the data +/// type of `y`. +/// +/// # Discussion +/// +/// [`group_values`] does not store any aggregate state inline. It only +/// assigns "group indices", one for each (distinct) group value. The +/// accumulators manage the in-progress aggregate state for each +/// group, with the group values themselves are stored in +/// [`group_values`] at the corresponding group index. +/// +/// The accumulator state (e.g partial sums) is managed by and stored +/// by a [`GroupsAccumulator`] accumulator. There is one accumulator +/// per aggregate expression (COUNT, AVG, etc) in the +/// stream. Internally, each `GroupsAccumulator` manages the state for +/// multiple groups, and is passed `group_indexes` during update. Note +/// The accumulator state is not managed by this operator (e.g in the +/// hash table). +/// +/// [`group_values`]: Self::group_values +/// +/// # Partial Aggregate and multi-phase grouping +/// +/// As described on [`Accumulator::state`], this operator is used in the context +/// "multi-phase" grouping when the mode is [`AggregateMode::Partial`]. +/// +/// An important optimization for multi-phase partial aggregation is to skip +/// partial aggregation when it is not effective enough to warrant the memory or +/// CPU cost, as is often the case for queries many distinct groups (high +/// cardinality group by). Memory is particularly important because each Partial +/// aggregator must store the intermediate state for each group. +/// +/// If the ratio of the number of groups to the number of input rows exceeds a +/// threshold, this operator will stop applying Partial aggregation and directly +/// pass the input rows to the next aggregation phase. +/// +/// [`Accumulator::state`]: datafusion_expr::Accumulator::state +/// +/// # Spilling (to disk) +/// +/// The sizes of group values and accumulators can become large. Before that causes out of memory, +/// this hash aggregator outputs partial states early for partial aggregation or spills to local +/// disk using Arrow IPC format for final aggregation. For every input [`RecordBatch`], the memory +/// manager checks whether the new input size meets the memory configuration. If not, outputting or +/// spilling happens. For outputting, the final aggregation takes care of re-grouping. For spilling, +/// later stream-merge sort on reading back the spilled data does re-grouping. Note the rows cannot +/// be grouped once spilled onto disk, the read back data needs to be re-grouped again. In addition, +/// re-grouping may cause out of memory again. Thus, re-grouping has to be a sort based aggregation. +/// ```text +/// Partial Aggregation [batch_size = 2] (max memory = 3 rows) +/// +/// INPUTS PARTIALLY AGGREGATED (UPDATE BATCH) OUTPUTS +/// ┌─────────┐ ┌─────────────────┐ ┌─────────────────┐ +/// │ a │ b │ │ a │ AVG(b) │ │ a │ AVG(b) │ +/// │---│-----│ │ │[count]│[sum]│ │ │[count]│[sum]│ +/// │ 3 │ 3.0 │ ─▶ │---│-------│-----│ │---│-------│-----│ +/// │ 2 │ 2.0 │ │ 2 │ 1 │ 2.0 │ ─▶ early emit ─▶ │ 2 │ 1 │ 2.0 │ +/// └─────────┘ │ 3 │ 2 │ 7.0 │ │ │ 3 │ 2 │ 7.0 │ +/// ┌─────────┐ ─▶ │ 4 │ 1 │ 8.0 │ │ └─────────────────┘ +/// │ 3 │ 4.0 │ └─────────────────┘ └▶ ┌─────────────────┐ +/// │ 4 │ 8.0 │ ┌─────────────────┐ │ 4 │ 1 │ 8.0 │ +/// └─────────┘ │ a │ AVG(b) │ ┌▶ │ 1 │ 1 │ 1.0 │ +/// ┌─────────┐ │---│-------│-----│ │ └─────────────────┘ +/// │ 1 │ 1.0 │ ─▶ │ 1 │ 1 │ 1.0 │ ─▶ early emit ─▶ ┌─────────────────┐ +/// │ 3 │ 2.0 │ │ 3 │ 1 │ 2.0 │ │ 3 │ 1 │ 2.0 │ +/// └─────────┘ └─────────────────┘ └─────────────────┘ +/// +/// +/// Final Aggregation [batch_size = 2] (max memory = 3 rows) +/// +/// PARTIALLY INPUTS FINAL AGGREGATION (MERGE BATCH) RE-GROUPED (SORTED) +/// ┌─────────────────┐ [keep using the partial schema] [Real final aggregation +/// │ a │ AVG(b) │ ┌─────────────────┐ output] +/// │ │[count]│[sum]│ │ a │ AVG(b) │ ┌────────────┐ +/// │---│-------│-----│ ─▶ │ │[count]│[sum]│ │ a │ AVG(b) │ +/// │ 3 │ 3 │ 3.0 │ │---│-------│-----│ ─▶ spill ─┐ │---│--------│ +/// │ 2 │ 2 │ 1.0 │ │ 2 │ 2 │ 1.0 │ │ │ 1 │ 4.0 │ +/// └─────────────────┘ │ 3 │ 4 │ 8.0 │ ▼ │ 2 │ 1.0 │ +/// ┌─────────────────┐ ─▶ │ 4 │ 1 │ 7.0 │ Streaming ─▶ └────────────┘ +/// │ 3 │ 1 │ 5.0 │ └─────────────────┘ merge sort ─▶ ┌────────────┐ +/// │ 4 │ 1 │ 7.0 │ ┌─────────────────┐ ▲ │ a │ AVG(b) │ +/// └─────────────────┘ │ a │ AVG(b) │ │ │---│--------│ +/// ┌─────────────────┐ │---│-------│-----│ ─▶ memory ─┘ │ 3 │ 2.0 │ +/// │ 1 │ 2 │ 8.0 │ ─▶ │ 1 │ 2 │ 8.0 │ │ 4 │ 7.0 │ +/// │ 2 │ 2 │ 3.0 │ │ 2 │ 2 │ 3.0 │ └────────────┘ +/// └─────────────────┘ └─────────────────┘ +/// ``` +pub(crate) struct GroupedHashAggregateStream { + // ======================================================================== + // PROPERTIES: + // These fields are initialized at the start and remain constant throughout + // the execution. + // ======================================================================== + schema: SchemaRef, + input_schema: SchemaRef, + input: SendableRecordBatchStream, + mode: AggregateMode, + + /// Arguments to pass to each accumulator. + /// + /// The arguments in `accumulator[i]` is passed `aggregate_arguments[i]` + /// + /// The argument to each accumulator is itself a `Vec` because + /// some aggregates such as `CORR` can accept more than one + /// argument. + aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>, + + /// Optional filter expression to evaluate, one for each for + /// accumulator. If present, only those rows for which the filter + /// evaluate to true should be included in the aggregate results. + /// + /// For example, for an aggregate like `SUM(x) FILTER (WHERE x >= 100)`, + /// the filter expression is `x > 100`. + filter_expressions: Arc<[Option<Arc<dyn PhysicalExpr>>]>, + + /// GROUP BY expressions + group_by: Arc<PhysicalGroupBy>, + + /// max rows in output RecordBatches + batch_size: usize, + + /// Optional soft limit on the number of `group_values` in a batch + /// If the number of `group_values` in a single batch exceeds this value, + /// the `GroupedHashAggregateStream` operation immediately switches to + /// output mode and emits all groups. + group_values_soft_limit: Option<usize>, + + // ======================================================================== + // STATE FLAGS: + // These fields will be updated during the execution. And control the flow of + // the execution. + // ======================================================================== + /// Tracks if this stream is generating input or output + exec_state: ExecutionState, + + /// Have we seen the end of the input + input_done: bool, + + // ======================================================================== + // STATE BUFFERS: + // These fields will accumulate intermediate results during the execution. + // ======================================================================== + /// An interning store of group keys + group_values: Box<dyn GroupValues>, + + /// scratch space for the current input [`RecordBatch`] being + /// processed. Reused across batches here to avoid reallocations + current_group_indices: Vec<usize>, + + /// Accumulators, one for each `AggregateFunctionExpr` in the query + /// + /// For example, if the query has aggregates, `SUM(x)`, + /// `COUNT(y)`, there will be two accumulators, each one + /// specialized for that particular aggregate and its input types + accumulators: Vec<Box<dyn GroupsAccumulator>>, + + // ======================================================================== + // TASK-SPECIFIC STATES: + // Inner states groups together properties, states for a specific task. + // ======================================================================== + /// Optional ordering information, that might allow groups to be + /// emitted from the hash table prior to seeing the end of the + /// input + group_ordering: GroupOrdering, + + /// The spill state object + spill_state: SpillState, + + /// Optional probe for skipping data aggregation, if supported by + /// current stream. + skip_aggregation_probe: Option<SkipAggregationProbe>, + + // ======================================================================== + // EXECUTION RESOURCES: + // Fields related to managing execution resources and monitoring performance. + // ======================================================================== + /// The memory reservation for this grouping + reservation: MemoryReservation, + + /// The behavior to trigger when out of memory occurs + oom_mode: OutOfMemoryMode, + + /// Execution metrics + baseline_metrics: BaselineMetrics, + + /// Aggregation-specific metrics + group_by_metrics: GroupByMetrics, + + /// Reduction factor metric, calculated as `output_rows/input_rows` (only for partial aggregation) + reduction_factor: Option<metrics::RatioMetrics>, +} + +impl GroupedHashAggregateStream { + /// Create a new GroupedHashAggregateStream + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug!("Creating GroupedHashAggregateStream"); + let agg_schema = Arc::clone(&agg.schema); + let agg_group_by = Arc::clone(&agg.group_by); + let agg_filter_expr = Arc::clone(&agg.filter_expr); + + let batch_size = context.session_config().batch_size(); + let input = agg.input.execute(partition, Arc::clone(context))?; + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + + let timer = baseline_metrics.elapsed_compute().timer(); + + let aggregate_exprs = Arc::clone(&agg.aggr_expr); + + // arguments for each aggregate, one vec of expressions per + // aggregate + let aggregate_arguments = aggregates::aggregate_expressions( + &agg.aggr_expr, + &agg.mode, + agg_group_by.num_group_exprs(), + )?; + // arguments for aggregating spilled data is the same as the one for final aggregation + let merging_aggregate_arguments = aggregates::aggregate_expressions( + &agg.aggr_expr, + &AggregateMode::Final, + agg_group_by.num_group_exprs(), + )?; + + let filter_expressions = match agg.mode.input_mode() { + AggregateInputMode::Raw => agg_filter_expr, + AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), + }; + + // Instantiate the accumulators + let accumulators: Vec<_> = aggregate_exprs + .iter() + .map(create_group_accumulator) + .collect::<Result<_>>()?; + + let group_schema = agg_group_by.group_schema(&agg.input().schema())?; + + // fix https://github.com/apache/datafusion/issues/13949 + // Builds a **partial aggregation** schema by combining the group columns and + // the accumulator state columns produced by each aggregate expression. + // + // # Why Partial Aggregation Schema Is Needed + // + // In a multi-stage (partial/final) aggregation strategy, each partial-aggregate + // operator produces *intermediate* states (e.g., partial sums, counts) rather + // than final scalar values. These extra columns do **not** exist in the original + // input schema (which may be something like `[colA, colB, ...]`). Instead, + // each aggregator adds its own internal state columns (e.g., `[acc_state_1, acc_state_2, ...]`). + // + // Therefore, when we spill these intermediate states or pass them to another + // aggregation operator, we must use a schema that includes both the group + // columns **and** the partial-state columns. + let spill_schema = Arc::new(create_schema( + &agg.input().schema(), + &agg_group_by, + &aggregate_exprs, + AggregateMode::Partial, + )?); + + // Need to update the GROUP BY expressions to point to the correct column after schema change + let merging_group_by_expr = agg_group_by + .expr + .iter() + .enumerate() + .map(|(idx, (_, name))| { + (Arc::new(Column::new(name.as_str(), idx)) as _, name.clone()) + }) + .collect(); + + let output_ordering = agg.cache.output_ordering(); + + let spill_sort_exprs = + group_schema + .fields + .into_iter() + .enumerate() + .map(|(idx, field)| { + let output_expr = Column::new(field.name().as_str(), idx); + + // Try to use the sort options from the output ordering, if available. + // This ensures that spilled state is sorted in the required order as well. + let sort_options = output_ordering + .and_then(|o| o.get_sort_options(&output_expr)) + .unwrap_or_default(); + + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_ordering) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Spill expression is empty"); + }; + + let agg_fn_names = aggregate_exprs + .iter() + .map(|expr| { + format_human_display(expr.human_display(), expr.human_display_alias()) + .map(|display| display.into_owned()) + .unwrap_or_else(|| expr.name().to_string()) + }) + .collect::<Vec<_>>() + .join(", "); + let name = format!("GroupedHashAggregateStream[{partition}] ({agg_fn_names})"); + let group_ordering = GroupOrdering::try_new(&agg.input_order_mode)?; + let oom_mode = match (agg.mode, &group_ordering) { + // In partial aggregation mode, always prefer to emit incomplete results early. + (AggregateMode::Partial, _) => OutOfMemoryMode::EmitEarly, + // For non-partial aggregation modes, emitting incomplete results is not an option. + // Instead, use disk spilling to store sorted, incomplete results, and merge them + // afterwards. + (_, GroupOrdering::None | GroupOrdering::Partial(_)) + if context.runtime_env().disk_manager.tmp_files_enabled() => + { + OutOfMemoryMode::Spill + } + // For `GroupOrdering::Full`, the incoming stream is already sorted. This ensures the + // number of incomplete groups can be kept small at all times. If we still hit + // an out-of-memory condition, spilling to disk would not be beneficial since the same + // situation is likely to reoccur when reading back the spilled data. + // Therefore, we fall back to simply reporting the error immediately. + // This mode will also be used if the `DiskManager` is not configured to allow spilling + // to disk. + _ => OutOfMemoryMode::ReportError, + }; + + let group_values = new_group_values(group_schema, &group_ordering)?; + let reservation = MemoryConsumer::new(name) + // We interpret 'can spill' as 'can handle memory back pressure'. + // This value needs to be set to true for the default memory pool implementations + // to ensure fair application of back pressure amongst the memory consumers. + .with_can_spill(oom_mode != OutOfMemoryMode::ReportError) + .register(context.memory_pool()); + timer.done(); + + let exec_state = ExecutionState::ReadingInput; + + let spill_manager = SpillManager::new( + context.runtime_env(), + metrics::SpillMetrics::new(&agg.metrics, partition), + Arc::clone(&spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + let spill_state = SpillState { + spills: vec![], + spill_expr: spill_ordering, + spill_schema, + is_stream_merging: false, + merging_aggregate_arguments, + merging_group_by: PhysicalGroupBy::new_single(merging_group_by_expr), + peak_mem_used: MetricBuilder::new(&agg.metrics) + .peak_memory_usage("peak_mem_used", partition), + spill_manager, + }; + + // Skip aggregation is supported if: + // - aggregation mode is Partial + // - input is not ordered by GROUP BY expressions, + // since Final mode expects unique group values as its input + // - there is only one GROUP BY expressions set + let skip_aggregation_probe = if agg.mode == AggregateMode::Partial + && matches!(group_ordering, GroupOrdering::None) + && agg_group_by.is_single() + { + let options = &context.session_config().options().execution; + let probe_rows_threshold = + options.skip_partial_aggregation_probe_rows_threshold; + let probe_ratio_threshold = + options.skip_partial_aggregation_probe_ratio_threshold; + // A threshold >= 1.0 means the ratio (num_groups / input_rows) can + // never exceed it, so the feature is effectively disabled. + if probe_ratio_threshold >= 1.0 { + None + } else { + let skipped_aggregation_rows = MetricBuilder::new(&agg.metrics) + .with_category(MetricCategory::Rows) + .counter("skipped_aggregation_rows", partition); + Some(SkipAggregationProbe::new( + probe_rows_threshold, + probe_ratio_threshold, + skipped_aggregation_rows, + )) + } + } else { + None + }; + + let reduction_factor = if agg.mode == AggregateMode::Partial { + Some( + MetricBuilder::new(&agg.metrics) + .with_type(metrics::MetricType::Summary) + .ratio_metrics("reduction_factor", partition), + ) + } else { + None + }; + + Ok(GroupedHashAggregateStream { + schema: agg_schema, + input_schema: agg.input().schema(), + input, + mode: agg.mode, + accumulators, + aggregate_arguments, + filter_expressions, + group_by: agg_group_by, + reservation, + oom_mode, + group_values, + current_group_indices: Default::default(), + exec_state, + baseline_metrics, + group_by_metrics, + batch_size, + group_ordering, + input_done: false, + spill_state, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + skip_aggregation_probe, + reduction_factor, + }) + } +} + +/// Create an accumulator for `agg_expr` -- a [`GroupsAccumulator`] if +/// that is supported by the aggregate, or a +/// [`GroupsAccumulatorAdapter`] if not. +pub(crate) fn create_group_accumulator( + agg_expr: &Arc<AggregateFunctionExpr>, +) -> Result<Box<dyn GroupsAccumulator>> { + if agg_expr.groups_accumulator_supported() { + agg_expr.create_groups_accumulator() + } else { + // Note in the log when the slow path is used + debug!( + "Creating GroupsAccumulatorAdapter for {}: {agg_expr:?}", + agg_expr.name() + ); + let agg_expr_captured = Arc::clone(agg_expr); + let factory = move || agg_expr_captured.create_accumulator(); + Ok(Box::new(GroupsAccumulatorAdapter::new(factory))) + } +} + +impl Stream for GroupedHashAggregateStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + + loop { + match &self.exec_state { + ExecutionState::ReadingInput => 'reading_input: { + match ready!(self.input.poll_next_unpin(cx)) { + // New batch to aggregate + Some(Ok(batch)) => { + let timer = elapsed_compute.timer(); + let input_rows = batch.num_rows(); + + if self.mode == AggregateMode::Partial + && let Some(reduction_factor) = + self.reduction_factor.as_ref() + { + reduction_factor.add_total(input_rows); + } + + // Do the grouping. + // `group_aggregate_batch` will _not_ have updated the memory reservation yet. + // The rest of the code will first try to reduce memory usage by + // already emitting results. + self.group_aggregate_batch(&batch)?; + + assert!(!self.input_done); + + // If the number of group values equals or exceeds the soft limit, + // emit all groups and switch to producing output + if self.hit_soft_group_limit() { + timer.done(); + self.set_input_done_and_produce_output()?; + // make sure the exec_state just set is not overwritten below + break 'reading_input; + } + + // Try to emit completed groups if possible. + // If we already started spilling, we can no longer emit since + // this might lead to incorrect output ordering + if (self.spill_state.spills.is_empty() + || self.spill_state.is_stream_merging) + && let Some(to_emit) = self.group_ordering.emit_to() + { + timer.done(); + if let Some(batch) = self.emit(to_emit, false)? { + self.exec_state = + ExecutionState::ProducingOutput(batch); + }; + // make sure the exec_state just set is not overwritten below + break 'reading_input; + } + + if self.mode == AggregateMode::Partial { + // Spilling should never be activated in partial aggregation mode. + assert!(!self.spill_state.is_stream_merging); + + // Check if we should switch to skip aggregation mode + // It's important that we do this before we early emit since we've + // already updated the probe. + self.update_skip_aggregation_probe(input_rows); + if let Some(new_state) = + self.switch_to_skip_aggregation()? + { + timer.done(); + self.exec_state = new_state; + break 'reading_input; + } + } + + // If we reach this point, try to update the memory reservation + // handling out-of-memory conditions as determined by the OOM mode. + if let Some(new_state) = + self.try_update_memory_reservation()? + { + timer.done(); + self.exec_state = new_state; + break 'reading_input; + } + + timer.done(); + } + + // Found error from input stream + Some(Err(e)) => { + // inner had error, return to caller + return Poll::Ready(Some(Err(e))); + } + + // Found end from input stream + None => { + // inner is done, emit all rows and switch to producing output + self.set_input_done_and_produce_output()?; + } + } + } + + ExecutionState::SkippingAggregation => { + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + let _timer = elapsed_compute.timer(); + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.record_skipped(&batch); + } + let states = self.transform_to_states(&batch)?; + return Poll::Ready(Some(Ok( + states.record_output(&self.baseline_metrics) + ))); + } + Some(Err(e)) => { + // inner had error, return to caller + return Poll::Ready(Some(Err(e))); + } + None => { + // inner is done, switching to `Done` state + // Sanity check: when switching from SkippingAggregation to Done, + // all groups should have already been emitted + if !self.group_values.is_empty() { + return Poll::Ready(Some(internal_err!( + "Switching from SkippingAggregation to Done with {} groups still in hash table. \ + This is a bug - all groups should have been emitted before skip aggregation started.", + self.group_values.len() + ))); + } + // Release the input pipeline's resources. + let input_schema = self.input.schema(); + self.input = + Box::pin(EmptyRecordBatchStream::new(input_schema)); + self.exec_state = ExecutionState::Done; + } + } + } + + ExecutionState::ProducingOutput(batch) => { + // slice off a part of the batch, if needed + let output_batch; + let size = self.batch_size; + (self.exec_state, output_batch) = if batch.num_rows() <= size { + ( + if self.input_done { + ExecutionState::Done + } + // In Partial aggregation, we also need to check + // if we should trigger partial skipping + else if self.mode == AggregateMode::Partial + && self.should_skip_aggregation() + { + ExecutionState::SkippingAggregation + } else { + ExecutionState::ReadingInput + }, + batch.clone(), + ) + } else { + // output first batch_size rows + let size = self.batch_size; + let num_remaining = batch.num_rows() - size; + let remaining = batch.slice(size, num_remaining); + let output = batch.slice(0, size); + (ExecutionState::ProducingOutput(remaining), output) + }; + + if let Some(reduction_factor) = self.reduction_factor.as_ref() { + reduction_factor.add_part(output_batch.num_rows()); + } + + // Empty record batches should not be emitted. + // They need to be treated as [`Option<RecordBatch>`]es and handled separately + debug_assert!(output_batch.num_rows() > 0); + return Poll::Ready(Some(Ok( + output_batch.record_output(&self.baseline_metrics) + ))); + } + + ExecutionState::Done => { + // Sanity check: all groups should have been emitted by now + if !self.group_values.is_empty() { + return Poll::Ready(Some(internal_err!( + "AggregateStream was in Done state with {} groups left in hash table. \ + This is a bug - all groups should have been emitted before entering Done state.", + self.group_values.len() + ))); + } + // release the memory reservation since sending back output batch itself needs + // some memory reservation, so make some room for it. + self.clear_all(); + let _ = self.update_memory_reservation(); + return Poll::Ready(None); + } + } + } + } +} + +impl RecordBatchStream for GroupedHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl GroupedHashAggregateStream { + /// Perform group-by aggregation for the given [`RecordBatch`]. + fn group_aggregate_batch(&mut self, batch: &RecordBatch) -> Result<()> { + // Evaluate the grouping expressions + let group_by_values = if self.spill_state.is_stream_merging { + evaluate_group_by(&self.spill_state.merging_group_by, batch)? + } else { + evaluate_group_by(&self.group_by, batch)? + }; + + // Only create the timer if there are actual aggregate arguments to evaluate + let timer = match ( + self.spill_state.is_stream_merging, + self.spill_state.merging_aggregate_arguments.is_empty(), + self.aggregate_arguments.is_empty(), + ) { + (true, false, _) | (false, _, false) => { + Some(self.group_by_metrics.aggregate_arguments_time.timer()) + } + _ => None, + }; + + // Evaluate the aggregation expressions. + let input_values = if self.spill_state.is_stream_merging { + evaluate_many(&self.spill_state.merging_aggregate_arguments, batch)? + } else { + evaluate_many(&self.aggregate_arguments, batch)? + }; + drop(timer); + + // Evaluate the filter expressions, if any, against the inputs + let filter_values = if self.spill_state.is_stream_merging { + let filter_expressions = vec![None; self.accumulators.len()]; + evaluate_optional(&filter_expressions, batch)? + } else { + evaluate_optional(&self.filter_expressions, batch)? + }; + + for group_values in &group_by_values { + let groups_start_time = Instant::now(); + + // calculate the group indices for each input row + let starting_num_groups = self.group_values.len(); + self.group_values + .intern(group_values, &mut self.current_group_indices)?; + let group_indices = &self.current_group_indices; + + // Update ordering information if necessary + let total_num_groups = self.group_values.len(); + if total_num_groups > starting_num_groups { + self.group_ordering.new_groups( + group_values, + group_indices, + total_num_groups, + )?; + } + + // Use this instant for both measurements to save a syscall + let agg_start_time = Instant::now(); + self.group_by_metrics + .time_calculating_group_ids + .add_duration(agg_start_time - groups_start_time); + + // Gather the inputs to call the actual accumulator + let t = self + .accumulators + .iter_mut() + .zip(input_values.iter()) + .zip(filter_values.iter()); + + for ((acc, values), opt_filter) in t { + let opt_filter = opt_filter.as_ref().map(|filter| filter.as_boolean()); + + // Call the appropriate method on each aggregator with + // the entire input row and the relevant group indexes + if self.mode.input_mode() == AggregateInputMode::Raw + && !self.spill_state.is_stream_merging + { + acc.update_batch( + values, + group_indices, + opt_filter, + total_num_groups, + )?; + } else { + assert_or_internal_err!( + opt_filter.is_none(), + "aggregate filter should be applied in partial stage, there should be no filter in final stage" + ); + + // if aggregation is over intermediate states, + // use merge + acc.merge_batch(values, group_indices, total_num_groups)?; + } + self.group_by_metrics + .aggregation_time + .add_elapsed(agg_start_time); + } + } + + Ok(()) + } + + /// Attempts to update the memory reservation. If that fails due to a + /// [DataFusionError::ResourcesExhausted] error, an attempt will be made to resolve + /// the out-of-memory condition based on the [out-of-memory handling mode](OutOfMemoryMode). + /// + /// If the out-of-memory condition can not be resolved, an `Err` value will be returned + /// + /// Returns `Ok(Some(ExecutionState))` if the state should be changed, `Ok(None)` otherwise. + fn try_update_memory_reservation(&mut self) -> Result<Option<ExecutionState>> { + let oom = match self.update_memory_reservation() { + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Err(e), + Ok(_) => return Ok(None), + }; + + match self.oom_mode { + OutOfMemoryMode::Spill if !self.group_values.is_empty() => { + self.spill()?; + self.clear_shrink(self.batch_size); + self.update_memory_reservation()?; + Ok(None) + } + OutOfMemoryMode::EmitEarly if self.group_values.len() > 1 => { + let n = if self.group_values.len() >= self.batch_size { + // Try to emit an integer multiple of batch size if possible + self.group_values.len() / self.batch_size * self.batch_size + } else { + // Otherwise emit whatever we can + self.group_values.len() + }; + + if let Some(emit_to) = self.group_ordering.oom_emit_to(n) + && let Some(batch) = self.emit(emit_to, false)? + { + return Ok(Some(ExecutionState::ProducingOutput(batch))); + } + Err(oom) + } + OutOfMemoryMode::EmitEarly + | OutOfMemoryMode::Spill + | OutOfMemoryMode::ReportError => Err(oom), + } + } + + fn update_memory_reservation(&mut self) -> Result<()> { + let acc = self.accumulators.iter().map(|x| x.size()).sum::<usize>(); + let groups_and_acc_size = acc + + self.group_values.size() + + self.group_ordering.size() + + self.current_group_indices.allocated_size(); + + // Reserve extra headroom for sorting during potential spill. + // When OOM triggers, group_aggregate_batch has already processed the + // latest input batch, so the internal state may have grown well beyond + // the last successful reservation. The emit batch reflects this larger + // actual state, and the sort needs memory proportional to it. + // By reserving headroom equal to the data size, we trigger OOM earlier + // (before too much data accumulates), ensuring the freed reservation + // after clear_shrink is sufficient to cover the sort memory. + let sort_headroom = + if self.oom_mode == OutOfMemoryMode::Spill && !self.group_values.is_empty() { + acc + self.group_values.size() + } else { + 0 + }; + + let new_size = groups_and_acc_size + sort_headroom; + let reservation_result = self.reservation.try_resize(new_size); + + if reservation_result.is_ok() { + self.spill_state + .peak_mem_used + .set_max(self.reservation.size()); + } + + reservation_result + } + + /// Create an output RecordBatch with the group keys and + /// accumulator states/values specified in emit_to + fn emit(&mut self, emit_to: EmitTo, spilling: bool) -> Result<Option<RecordBatch>> { + let schema = if spilling { + Arc::clone(&self.spill_state.spill_schema) + } else { + self.schema() + }; + if self.group_values.is_empty() { + return Ok(None); + } + + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.group_values.emit(emit_to)?; + if let EmitTo::First(n) = emit_to { + self.group_ordering.remove_groups(n); + } + + // Next output each aggregate value + for acc in self.accumulators.iter_mut() { + if self.mode.output_mode() == AggregateOutputMode::Final && !spilling { + output.push(acc.evaluate(emit_to)?) + } else { + // Output partial state: either because we're in a non-final mode, + // or because we're spilling and will merge/re-evaluate later. + output.extend(acc.state(emit_to)?) + } + } + drop(timer); + + // emit reduces the memory usage. Ignore Err from update_memory_reservation. Even if it is + // over the target memory size after emission, we can emit again rather than returning Err. + let _ = self.update_memory_reservation(); + let batch = RecordBatch::try_new(schema, output)?; + debug_assert!(batch.num_rows() > 0); + + Ok(Some(batch)) + } + + /// Registers groups for empty grouping sets when no input rows were seen. + /// + /// `GROUP BY GROUPING SETS (())` must always produce one row even when there + /// are no input rows (standard SQL semantics for a "grand total" group). + /// Mixed grouping sets like `GROUPING SETS (a, ())` also produce one row for + /// the empty set `()` on empty input. + /// + /// This method interns the group keys and primes the accumulators so they + /// produce their zero-row aggregate values (e.g. `NULL` for `SUM`, + /// `0` for `COUNT`). + fn init_empty_grouping_sets(&mut self) -> Result<()> { + if !self.group_by.has_grouping_set() || !self.group_values.is_empty() { + return Ok(()); + } + + let max_ordinal = max_duplicate_ordinal(self.group_by.groups()); + let mut ordinals: std::collections::HashMap<&[bool], usize> = + std::collections::HashMap::new(); + let group_schema = self.group_by.group_schema(&self.input_schema)?; + let n_expr = self.group_by.expr().len(); + let mut any_interned = false; + + for group in self.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let o = *entry; + *entry += 1; + o + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + // Build the group key: one NULL per group-by expression, then the grouping_id. + let mut cols: Vec<ArrayRef> = group_schema + .fields() + .iter() + .take(n_expr) + .map(|f| new_null_array(f.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + let starting_groups = self.group_values.len(); + self.group_values + .intern(&cols, &mut self.current_group_indices)?; + let total_groups = self.group_values.len(); + if total_groups > starting_groups { + self.group_ordering.new_groups( + &cols, + &self.current_group_indices, + total_groups, + )?; + } + any_interned = true; + } + + if any_interned { + // Prime each accumulator for the registered group count with no data. + // + // We build 1-row null arrays for each aggregate argument and pass them + // with an all-false filter to update_batch. The filter ensures no row + // is accumulated into any group, which keeps every group in its "zero" + // initial state (NULL for SUM/AVG/MIN/MAX, 0 for COUNT). + // + // Using a 1-row batch rather than 0 rows is required to avoid a fast + // path in `NullState::accumulate` that treats "0 nulls in a 0-row + // array" as "all groups have been seen", which would cause SUM to + // return 0 instead of NULL. + // + // This path always runs in a Raw input mode, so `update_batch` (not + // `merge_batch`) is the right entry point: + // + // - `has_grouping_set()` can only be true for the Partial / Single / + // SinglePartitioned modes, whose `input_mode()` is `Raw`. The final + // modes rebuild their group-by via `PhysicalGroupBy::as_final()`, + // which clears `has_grouping_set`, so this method returns early for + // them and never reaches here. + // + // Since every row is filtered out, the actual data content never + // matters. The assert documents and guards the invariant above. + debug_assert_eq!( + self.mode.input_mode(), + AggregateInputMode::Raw, + "init_empty_grouping_sets must only run in a Raw input mode" + ); + let total_groups = self.group_values.len(); + let null_args: Vec<Vec<ArrayRef>> = self + .aggregate_arguments + .iter() + .map(|args| { + args.iter() + .map(|expr| { + let dt = expr.data_type(&self.input_schema)?; + Ok(new_null_array(&dt, 1)) + }) + .collect::<Result<Vec<_>>>() + }) + .collect::<Result<Vec<_>>>()?; + let false_filter = BooleanArray::from(vec![false]); + for (acc, args) in self.accumulators.iter_mut().zip(null_args.iter()) { + acc.update_batch(args, &[0], Some(&false_filter), total_groups)?; + } + } + + Ok(()) + } + + /// Emit all intermediate aggregation states, sort them, and store them on disk. + /// This process helps in reducing memory pressure by allowing the data to be + /// read back with streaming merge. + fn spill(&mut self) -> Result<()> { + // Emit and sort intermediate aggregation state + let Some(emit) = self.emit(EmitTo::All, true)? else { + return Ok(()); + }; + + // Free accumulated state now that data has been emitted into `emit`. + // This must happen before reserving sort memory so the pool has room. + // Use 0 to minimize allocated capacity and maximize memory available for sorting. + self.clear_shrink(0); + self.update_memory_reservation()?; + + let batch_size_ratio = self.batch_size as f32 / emit.num_rows() as f32; + let batch_memory = get_record_batch_memory_size(&emit); + // The maximum worst case for a sort is 2X the original underlying buffers(regardless of slicing) + // First we get the underlying buffers' size, then we get the sliced("actual") size of the batch, + // and multiply it by the ratio of batch_size to actual size to get the estimated memory needed for sorting the batch. + // If something goes wrong in get_sliced_size()(double counting or something), + // we fall back to the worst case. + let sort_memory = (batch_memory + + (emit.get_sliced_size()? as f32 * batch_size_ratio) as usize) + .min(batch_memory * 2); + + // If we can't grow even that, we have no choice but to return an error since we can't spill to disk without sorting the data first. + self.reservation.try_grow(sort_memory).map_err(|err| { + resources_datafusion_err!( + "Failed to reserve memory for sort during spill: {err}" + ) + })?; + + let sorted_iter = IncrementalSortIterator::new( + emit, + self.spill_state.spill_expr.clone(), + self.batch_size, + ); + let spillfile = self + .spill_state + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "HashAggSpill", + )?; + + // Shrink the memory we allocated for sorting as the sorting is fully done at this point. + self.reservation.shrink(sort_memory); + + match spillfile { + Some((spillfile, max_record_batch_memory)) => { + self.spill_state.spills.push(SortedSpillFile { + file: spillfile, + max_record_batch_memory, + }) + } + None => { + return internal_err!( + "Calling spill with no intermediate batch to spill" + ); + } + } + + Ok(()) + } + + /// Clear memory and shrink capacities to the given number of rows. + fn clear_shrink(&mut self, num_rows: usize) { + self.group_values.clear_shrink(num_rows); + self.current_group_indices.clear(); + self.current_group_indices.shrink_to(num_rows); + } + + /// Clear memory and shrink capacities to zero. + fn clear_all(&mut self) { + self.clear_shrink(0); + } + + /// returns true if there is a soft groups limit and the number of distinct + /// groups we have seen is over that limit + fn hit_soft_group_limit(&self) -> bool { + let Some(group_values_soft_limit) = self.group_values_soft_limit else { + return false; + }; + group_values_soft_limit <= self.group_values.len() + } + + /// Finalizes reading of the input stream and prepares for producing output values. + /// + /// This method is called both when the original input stream and, + /// in case of disk spilling, the SPM stream have been drained. + fn set_input_done_and_produce_output(&mut self) -> Result<()> { + self.input_done = true; + self.group_ordering.input_done(); + // Release the original input pipeline's resources now that we're done + // reading from it. In the spill branch below, `self.input` is replaced + // again with a stream that merges spill files. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + self.exec_state = if self.spill_state.spills.is_empty() { + // Input has been entirely processed without spilling to disk. + self.init_empty_grouping_sets()?; + + // Flush any remaining group values. + let batch = self.emit(EmitTo::All, false)?; + + // If there are none, we're done; otherwise switch to emitting them + batch.map_or(ExecutionState::Done, ExecutionState::ProducingOutput) + } else { + // Spill any remaining data to disk. There is some performance overhead in + // writing out this last chunk of data and reading it back. The benefit of + // doing this is that memory usage for this stream is reduced, and the more + // sophisticated memory handling in `MultiLevelMergeBuilder` can take over + // instead. + // Spilling to disk and reading back also ensures batch size is consistent + // rather than potentially having one significantly larger last batch. + self.spill()?; + + // Mark that we're switching to stream merging mode. + self.spill_state.is_stream_merging = true; + + self.input = StreamingMergeBuilder::new() + .with_schema(Arc::clone(&self.spill_state.spill_schema)) + .with_spill_manager(self.spill_state.spill_manager.clone()) + .with_sorted_spill_files(std::mem::take(&mut self.spill_state.spills)) + .with_expressions(&self.spill_state.spill_expr) + .with_metrics(self.baseline_metrics.clone()) + .with_batch_size(self.batch_size) + .with_reservation(self.reservation.new_empty()) + .build()?; + self.input_done = false; + + // Reset the group values collectors. + self.clear_all(); + + // We can now use `GroupOrdering::Full` since the spill files are sorted + // on the grouping columns. + self.group_ordering = GroupOrdering::Full(GroupOrderingFull::new()); + + // Recreate `group_values` for streaming merge so group ids are assigned + // in first-seen order, as required by `GroupOrderingFull`. + // The pre-spill multi-column collector may use `vectorized_intern`, which + // can assign new group ids out of input order under hash collisions. + let group_schema = self + .spill_state + .merging_group_by + .group_schema(&self.spill_state.spill_schema)?; + if group_schema.fields().len() > 1 { + self.group_values = new_group_values(group_schema, &self.group_ordering)?; + } + + // Use `OutOfMemoryMode::ReportError` from this point on + // to ensure we don't spill the spilled data to disk again. + self.oom_mode = OutOfMemoryMode::ReportError; + + self.update_memory_reservation()?; + + ExecutionState::ReadingInput + }; + timer.done(); + Ok(()) + } + + /// Updates skip aggregation probe state. + /// + /// Notice: It should only be called in Partial aggregation + fn update_skip_aggregation_probe(&mut self, input_rows: usize) { + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + // Skip aggregation probe is not supported if stream has any spills, + // currently spilling is not supported for Partial aggregation + assert!(self.spill_state.spills.is_empty()); + probe.update_state(input_rows, self.group_values.len()); + }; + } + + /// In case the probe indicates that aggregation may be + /// skipped, forces stream to produce currently accumulated output. + /// + /// Notice: It should only be called in Partial aggregation + /// + /// Returns `Some(ExecutionState)` if the state should be changed, None otherwise. + fn switch_to_skip_aggregation(&mut self) -> Result<Option<ExecutionState>> { + if let Some(probe) = self.skip_aggregation_probe.as_mut() + && probe.should_skip() + && let Some(batch) = self.emit(EmitTo::All, false)? + { + return Ok(Some(ExecutionState::ProducingOutput(batch))); + }; + + Ok(None) + } + + /// Returns true if the aggregation probe indicates that aggregation + /// should be skipped. + /// + /// Notice: It should only be called in Partial aggregation + fn should_skip_aggregation(&self) -> bool { + self.skip_aggregation_probe + .as_ref() + .is_some_and(|probe| probe.should_skip()) + } + + /// Transforms input batch to intermediate aggregate state, without grouping it + fn transform_to_states(&self, batch: &RecordBatch) -> Result<RecordBatch> { + let mut group_values = evaluate_group_by(&self.group_by, batch)?; + let input_values = evaluate_many(&self.aggregate_arguments, batch)?; + let filter_values = evaluate_optional(&self.filter_expressions, batch)?; + + assert_eq_or_internal_err!( + group_values.len(), + 1, + "group_values expected to have single element" + ); + let mut output = group_values.swap_remove(0); + + let iter = self + .accumulators + .iter() + .zip(input_values.iter()) + .zip(filter_values.iter()); + + for ((acc, values), opt_filter) in iter { + let opt_filter = opt_filter.as_ref().map(|filter| filter.as_boolean()); + output.extend(acc.convert_to_state(values, opt_filter)?); + } + + let states_batch = RecordBatch::try_new(self.schema(), output)?; + + Ok(states_batch) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::InputOrderMode; + use crate::test::TestMemoryExec; + use arrow::array::{Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_double_emission_race_condition_bug() -> Result<()> { + // Fix for https://github.com/apache/datafusion/issues/18701 + // This test specifically proves that we have fixed double emission race condition + // where emit_early_if_necessary() and switch_to_skip_aggregation() + // both emit in the same loop iteration, causing data loss + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // Create data that will trigger BOTH conditions in the same iteration: + // 1. More groups than batch_size (triggers early emission when memory pressure hits) + // 2. High cardinality ratio (triggers skip aggregation) + let batch_size = 1024; // We'll set this in session config + let num_groups = batch_size + 100; // Slightly more than batch_size (1124 groups) + + // Create exactly 1 row per group = 100% cardinality ratio + let group_ids: Vec<i32> = (0..num_groups as i32).collect(); + let values: Vec<i64> = vec![1; num_groups]; + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids)), + Arc::new(Int64Array::from(values)), + ], + )?; + + let input_partitions = vec![vec![batch]]; + + // Create constrained memory to trigger early emission but not completely fail + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1024, 1.0) // small enough to start but will trigger pressure + .build_arc()?; + + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure to trigger BOTH conditions: + // 1. Low probe threshold (triggers skip probe after few rows) + // 2. Low ratio threshold (triggers skip aggregation immediately) + // 3. Set batch_size to 1024 so our 1124 groups will trigger early emission + // This creates the race condition where both emit paths are triggered + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.batch_size", + &datafusion_common::ScalarValue::UInt64(Some(1024)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(50)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(0.8)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode where the race condition occurs + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + GroupedHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Count total groups emitted + let mut total_output_groups = 0; + for batch in &results { + total_output_groups += batch.num_rows(); + } + + assert_eq!( + total_output_groups, num_groups, + "Unexpected number of groups", + ); + + Ok(()) + } + + // Migrated to OrderedPartialAggregateStream coverage in aggregates/mod.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_emit_early_with_partially_sorted() -> Result<()> { + // Reproducer for #20445: EmitEarly with PartiallySorted panics in + // remove_groups because it emits more groups than the sort boundary. + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // All rows share sort_col=1 (no sort boundary), with unique group_col + // values to create many groups and trigger memory pressure. + let n = 256; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; n])), + Arc::new(Int32Array::from((0..n as i32).collect::<Vec<_>>())), + Arc::new(Int64Array::from(vec![1; n])), + ], + )?; + + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(4096, 1.0) + .build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(runtime); + let mut cfg = task_ctx.session_config().clone(); + cfg = cfg.set( + "datafusion.execution.batch_size", + &datafusion_common::ScalarValue::UInt64(Some(128)), + ); + cfg = cfg.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(u64::MAX)), + ); + task_ctx = task_ctx.with_session_config(cfg); + let task_ctx = Arc::new(task_ctx); + + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ) + as _)]) + .unwrap(); + let exec = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // GROUP BY sort_col, group_col with input sorted on sort_col + // gives PartiallySorted([0]) + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )], + vec![None], + exec, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate_exec.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + // Must not panic with "assertion failed: *current_sort >= n" + let mut stream = GroupedHashAggregateStream::new(&aggregate_exec, &task_ctx, 0)?; + while let Some(result) = stream.next().await { + if let Err(e) = result { + if e.to_string().contains("Resources exhausted") { + break; + } + return Err(e); + } + } + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/grouped_topk_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/grouped_topk_stream.rs new file mode 100644 index 0000000..193fdba --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/grouped_topk_stream.rs
@@ -0,0 +1,306 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A memory-conscious aggregation implementation that limits group buckets to a fixed number + +use crate::aggregates::group_values::GroupByMetrics; +use crate::aggregates::topk::priority_map::PriorityMap; +#[cfg(debug_assertions)] +use crate::aggregates::topk_types_supported; +use crate::aggregates::{ + AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, + evaluate_many, +}; +use crate::metrics::BaselineMetrics; +use crate::stream::EmptyRecordBatchStream; +use crate::{RecordBatchStream, SendableRecordBatchStream}; +use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::compute::concat; +use arrow::datatypes::SchemaRef; +use arrow::util::pretty::print_batches; +use datafusion_common::Result; +use datafusion_common::internal_datafusion_err; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::metrics::RecordOutput; +use futures::stream::{Stream, StreamExt}; +use log::{Level, trace}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +pub struct GroupedTopKAggregateStream { + partition: usize, + row_count: usize, + started: bool, + done: bool, + schema: SchemaRef, + input: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, + group_by_metrics: GroupByMetrics, + aggregate_arguments: Vec<Vec<Arc<dyn PhysicalExpr>>>, + group_by: Arc<PhysicalGroupBy>, + priority_map: PriorityMap, + /// Whether a NULL group key has been seen for a group-by-only aggregation. + null_group_seen: bool, +} + +impl GroupedTopKAggregateStream { + pub fn new( + aggr: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + limit: usize, + ) -> Result<Self> { + let agg_schema = Arc::clone(&aggr.schema); + let group_by = Arc::clone(&aggr.group_by); + let input = aggr.input.execute(partition, Arc::clone(context))?; + let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition); + let group_by_metrics = GroupByMetrics::new(&aggr.metrics, partition); + let aggregate_arguments = + aggregate_expressions(&aggr.aggr_expr, &aggr.mode, group_by.expr.len())?; + + let (expr, _) = &aggr.group_expr().expr()[0]; + let kt = expr.data_type(&aggr.input().schema())?; + + // Check if this is a MIN/MAX aggregate or a DISTINCT-like operation + let (vt, desc) = if let Some((val_field, desc)) = aggr.get_minmax_desc() { + // MIN/MAX case: use the aggregate output type + (val_field.data_type().clone(), desc) + } else { + // DISTINCT case: use the group key type and get ordering from limit_order_descending + // The ordering direction is set by the optimizer when it pushes down the limit + let desc = aggr + .limit_options() + .and_then(|config| config.descending) + .ok_or_else(|| { + internal_datafusion_err!( + "Ordering direction required for DISTINCT with limit" + ) + })?; + (kt.clone(), desc) + }; + + // Type validation is performed by the optimizer and can_use_topk() check. + // This debug assertion documents the contract without runtime overhead in release builds. + #[cfg(debug_assertions)] + { + debug_assert!( + topk_types_supported(&kt, &vt), + "TopK type validation should have been performed by optimizer and can_use_topk(). \ + Found unsupported types: key={kt:?}, value={vt:?}" + ); + } + + // Note: Null values in aggregate columns are filtered by the aggregation layer + // before reaching the heap, so the heap implementations don't need explicit null handling. + let priority_map = PriorityMap::new(kt, vt, limit, desc)?; + + Ok(GroupedTopKAggregateStream { + partition, + started: false, + done: false, + row_count: 0, + schema: agg_schema, + input, + baseline_metrics, + group_by_metrics, + aggregate_arguments, + group_by, + priority_map, + null_group_seen: false, + }) + } +} + +impl RecordBatchStream for GroupedTopKAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl GroupedTopKAggregateStream { + fn is_group_by_only(&self) -> bool { + self.aggregate_arguments.is_empty() + } + + fn intern(&mut self, ids: &ArrayRef, vals: &ArrayRef) -> Result<()> { + let _timer = self.group_by_metrics.time_calculating_group_ids.timer(); + + let len = ids.len(); + self.priority_map + .set_batch(Arc::clone(ids), Arc::clone(vals)); + + let has_nulls = vals.null_count() > 0; + if has_nulls && self.is_group_by_only() { + self.null_group_seen = true; + } + // Keep the common no-NULL path free of NULL bookkeeping. Once a NULL + // group exists, use the NULL-aware path until it has been resolved. + let track_null_groups = !self.is_group_by_only() + && (has_nulls || self.priority_map.has_null_groups()); + for row_idx in 0..len { + if has_nulls && vals.is_null(row_idx) { + // MIN/MAX ignore NULL inputs, but a group whose values are all + // NULL must still be emitted with a NULL aggregate value, so + // track it. (GROUP BY-only aggregations handle NULL group keys + // via `null_group_seen` instead.) + if !self.is_group_by_only() { + self.priority_map.insert_null(row_idx); + } + continue; + } + if track_null_groups { + self.priority_map.insert_with_null_groups(row_idx)?; + } else { + self.priority_map.insert(row_idx)?; + } + } + Ok(()) + } + + fn emit_columns(&mut self) -> Result<Vec<ArrayRef>> { + let mut cols = if self.priority_map.is_empty() { + vec![] + } else { + self.priority_map.emit()? + }; + + // GROUP BY-only aggregation covers DISTINCT-like queries. The group + // key and heap value are the same column, but the output schema has + // only the group key. + if self.is_group_by_only() { + cols.truncate(1); + if self.null_group_seen { + self.append_null_group(&mut cols)?; + } + } + + Ok(cols) + } + + fn append_null_group(&self, cols: &mut Vec<ArrayRef>) -> Result<()> { + let dt = self.schema.field(0).data_type(); + let null_arr = new_null_array(dt, 1); + if cols.is_empty() { + cols.push(null_arr); + } else { + // NULL group keys are tracked outside the heap, so append a + // one-row NULL array to the emitted non-NULL group key column. + cols[0] = concat(&[cols[0].as_ref(), null_arr.as_ref()])?; + } + Ok(()) + } +} + +impl Stream for GroupedTopKAggregateStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.done { + return Poll::Ready(None); + } + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let emitting_time = self.group_by_metrics.emitting_time.clone(); + while let Poll::Ready(res) = self.input.poll_next_unpin(cx) { + let _timer = elapsed_compute.timer(); + match res { + // got a batch, convert to rows and append to our TreeMap + Some(Ok(batch)) => { + self.started = true; + trace!( + "partition {} has {} rows and got batch with {} rows", + self.partition, + self.row_count, + batch.num_rows() + ); + if log::log_enabled!(Level::Trace) && batch.num_rows() < 20 { + print_batches(std::slice::from_ref(&batch))?; + } + self.row_count += batch.num_rows(); + let batches = &[batch]; + let group_by_values = + evaluate_group_by(&self.group_by, batches.first().unwrap())?; + assert_eq!( + group_by_values.len(), + 1, + "Exactly 1 group value required" + ); + assert_eq!( + group_by_values[0].len(), + 1, + "Exactly 1 group value required" + ); + let group_by_values = Arc::clone(&group_by_values[0][0]); + let input_values = if self.is_group_by_only() { + // GROUP BY-only case: use group key as both key and value + Arc::clone(&group_by_values) + } else { + // MIN/MAX case: evaluate aggregate expressions + let _timer = + self.group_by_metrics.aggregate_arguments_time.timer(); + let input_values = evaluate_many( + &self.aggregate_arguments, + batches.first().unwrap(), + )?; + assert_eq!(input_values.len(), 1, "Exactly 1 input required"); + assert_eq!(input_values[0].len(), 1, "Exactly 1 input required"); + Arc::clone(&input_values[0][0]) + }; + + // iterate over each column of group_by values + (*self).intern(&group_by_values, &input_values)?; + } + // inner is done, emit all rows and switch to producing output + None => { + // Release the input pipeline's resources before emitting. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + if self.priority_map.is_empty() && !self.null_group_seen { + trace!("partition {} emit None", self.partition); + self.done = true; + return Poll::Ready(None); + } + let batch = { + let _timer = emitting_time.timer(); + let cols = self.emit_columns()?; + RecordBatch::try_new(Arc::clone(&self.schema), cols)? + }; + let batch = batch.record_output(&self.baseline_metrics); + trace!( + "partition {} emit batch with {} rows", + self.partition, + batch.num_rows() + ); + if log::log_enabled!(Level::Trace) { + print_batches(std::slice::from_ref(&batch))?; + } + self.done = true; + return Poll::Ready(Some(Ok(batch))); + } + // inner had error, return to caller + Some(Err(e)) => { + return Poll::Ready(Some(Err(e))); + } + } + } + Poll::Pending + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs new file mode 100644 index 0000000..f697e5a --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/hash_stream.rs
@@ -0,0 +1,1838 @@ +// 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. + +//! 2-stage hash aggregation stream implementation. +//! +//! See comments in [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] +//! for details. +//! +//! Note these streams are an incremental migration of the existing +//! [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! See issue for details: <https://github.com/apache/datafusion/issues/22710> + +use std::mem::size_of; +use std::ops::ControlFlow; +use std::sync::Arc; +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_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{ + AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, +}; +use super::group_values::GroupByMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::skip_partial::SkipAggregationProbe; +use crate::metrics::{ + BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, +}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; + +/// Hash aggregation is implemented in two stages: partial and final. This +/// stream implements the partial stage. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=final) +/// -- RepartitionExec(hash(k)) +/// ---- AggregateExec(stage=partial) +/// +/// ## Partial Stage Behavior +/// Input: raw rows +/// Output: partial states for all groups (for example, `AVG(x)` emits `SUM(x)` +/// and `COUNT(x)`) +/// +/// ## Final Stage Behavior +/// Input: partial states +/// Output: results for all groups (for example, `AVG(x)` calculated from the +/// state) +/// +/// # Optimization: DISTINCT LIMIT Soft Limit +/// +/// This optimization applies to both [`PartialHashAggregateStream`] and +/// [`FinalHashAggregateStream`]. +/// +/// Unordered distinct queries such as: +/// +/// ```sql +/// SELECT DISTINCT x FROM t LIMIT 10; +/// ``` +/// +/// are optimized into a two-stage aggregate like: +/// +/// ```txt +/// LimitExec, limit=10 +/// --AggregateExec(Final), group_by=[x], aggr=[], soft_limit=10 +/// ---- RepartitionExec, partitioning=hash(x) +/// ------ AggregateExec(Partial), group_by=[x], aggr=[], soft_limit=10 +/// -------- Scan(t) +/// ``` +/// +/// After each input batch, the stream checks whether the soft limit has been +/// reached. If so, it emits the accumulated groups and stops reading input. +/// +/// This operator does not guarantee an exact limit because a single batch can +/// cross the threshold. The downstream limit operator enforces the exact result +/// size. +/// +/// # Optimization: Partial Aggregation Skip +/// +/// Partial aggregation can be counterproductive for high-cardinality inputs, +/// where most rows create distinct groups. The stream probes the ratio of +/// accumulated groups to input rows while it is still aggregating. If the ratio +/// crosses the configured threshold and all aggregate accumulators can convert +/// raw inputs directly to partial state, the stream emits any already +/// accumulated groups, then switches to a skip state. In that state, each +/// remaining input batch is converted directly to partial aggregate state rows +/// without inserting the rows into the grouped hash table. +/// +/// # Feature: Memory-limited Execution +/// +/// ## Partial Aggregation +/// +/// Partial aggregation can emit incomplete results because the final stage merges +/// all intermediate states for the same group. If the memory reservation exceeds +/// its limit after aggregating an input batch, this stream emits all accumulated +/// states and continues aggregating the remaining input with an empty table. +/// +/// ## Final Aggregation +/// +/// During final aggregation, group keys and states accumulate. If memory usage +/// exceeds the budget, spilling is triggered as follows: +/// 1. After aggregating a new input batch, if the memory reservation exceeds its +/// limit, spill all accumulated groups and states. +/// - Sort all groups by the group keys before spilling. +/// 2. Repeat until the input is exhausted. +/// 3. Perform a sort-preserving merge of all spill files and feed the merged output +/// into an ordered streaming aggregation, which ensures bounded memory usage and +/// evaluates the final result. +/// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +pub(crate) struct PartialHashAggregateStream { + /// Output schema: group columns followed by partial aggregate state columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + input: SendableRecordBatchStream, + + /// Target output batch size from configuration. + batch_size: usize, + + /// Memory reservation for group keys and accumulators. + reservation: MemoryReservation, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Tracks partial aggregation row reduction, matching `GroupedHashAggregateStream`. + reduction_factor: metrics::RatioMetrics, + + /// Tracks whether partial aggregation should switch to direct state conversion. + skip_aggregation_probe: Option<SkipAggregationProbe>, + + /// Optional soft limit on the number of groups to accumulate before output. + /// + /// Invariant: when this is `Some(..)`, the accumulators inside `hash_table` must + /// be empty. See struct comments for details. + group_values_soft_limit: Option<usize>, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option<PartialHashAggregateState>, +} + +/// States for partial hash aggregation processing. +enum PartialHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable<PartialMarker>, + }, + /// A fully materialized partial-state batch being emitted incrementally. + EmittingOnMemoryPressure { + hash_table: AggregateHashTable<PartialMarker>, + // After each incremental emitting step, the `remaining_groups` will be updated + // with batch slicing. + remaining_groups: RecordBatch, + }, + ProducingOutput { + hash_table: AggregateHashTable<PartialMarker>, + /// If `None`, partial skip was never triggered and this state will + /// finish in `Done`. If `Some`, partial skip has triggered and the + /// stream will move to `SkippingAggregation` after these accumulated + /// groups are emitted. + skip_hash_table: Option<AggregateHashTable<PartialSkipMarker>>, + }, + SkippingAggregation { + hash_table: AggregateHashTable<PartialSkipMarker>, + }, + 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 PartialHashAggregatePoll = Poll<Option<Result<RecordBatch>>>; +type PartialHashAggregateStateTransition = ControlFlow< + (PartialHashAggregatePoll, PartialHashAggregateState), + PartialHashAggregateState, +>; + +/// Spill configuration and accumulated runs for final hash aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct FinalSpillContext { + /// Aggregate configuration used to construct the final replay stream. + final_agg: AggregateExec, + /// Task context. + context: Arc<TaskContext>, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Full group-key ordering kept by every spill file and the merged input. + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Spill runs waiting to be merged, they're all sorted by full group-by keys. + spills: Vec<SortedSpillFile>, +} + +/// Hash aggregation is implemented in two stages: partial and final. This +/// stream implements the final stage. +/// +/// See [`PartialHashAggregateStream`] for details. +pub(crate) struct FinalHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing partial aggregate state rows. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Memory reservation for group keys, accumulators, and spill sorting. + reservation: MemoryReservation, + + /// See comments for the same variable in [`PartialHashAggregateStream`]. + group_values_soft_limit: Option<usize>, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option<FinalHashAggregateState>, +} + +/// 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<FinalMarker>, + /// `None` if spilling is not supported by the configured `DiskManager`. + spill_context: Option<Box<FinalSpillContext>>, + }, + Spilling { + hash_table: AggregateHashTable<FinalMarker>, + spill_context: Box<FinalSpillContext>, + }, + ProducingOutput { + hash_table: AggregateHashTable<FinalMarker>, + }, + PreparingMergeInput { + hash_table: AggregateHashTable<FinalMarker>, + spill_context: Box<FinalSpillContext>, + }, + 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, +} + +type FinalHashAggregatePoll = Poll<Option<Result<RecordBatch>>>; +type FinalHashAggregateStateTransition = ControlFlow< + (FinalHashAggregatePoll, FinalHashAggregateState), + FinalHashAggregateState, +>; + +impl FinalSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + batch_size: usize, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result<Self> { + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + let output_ordering = agg.cache.output_ordering(); + let spill_sort_exprs = + group_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Final hash aggregate spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + let mut final_agg = agg.clone(); + final_agg.input_order_mode = InputOrderMode::Sorted; + + Ok(Self { + final_agg, + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`FinalHashAggregateStream`] for spilling details. + fn spill_table( + &mut self, + hash_table: &mut AggregateHashTable<FinalMarker>, + ) -> Result<()> { + let Some(batch) = hash_table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "FinalHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Final hash aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run, and do the aggregate evaluation with + /// [`OrderedFinalAggregateStream`] + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result<SendableRecordBatchStream> { + let Self { + final_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &final_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl PartialHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug_assert_eq!(agg.mode, super::AggregateMode::Partial); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reduction_factor = MetricBuilder::new(&agg.metrics) + .with_type(metrics::MetricType::Summary) + .ratio_metrics("reduction_factor", partition); + + let hash_table = AggregateHashTable::<PartialMarker>::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + let skip_aggregation_probe = if agg.group_by.is_single() { + let options = &context.session_config().options().execution; + let probe_ratio_threshold = + options.skip_partial_aggregation_probe_ratio_threshold; + // A threshold >= 1.0 means the ratio (num_groups / input_rows) can + // never exceed it, so the feature is effectively disabled. + if probe_ratio_threshold >= 1.0 { + None + } else { + let skipped_aggregation_rows = MetricBuilder::new(&agg.metrics) + .with_category(MetricCategory::Rows) + .counter("skipped_aggregation_rows", partition); + Some(SkipAggregationProbe::new( + options.skip_partial_aggregation_probe_rows_threshold, + probe_ratio_threshold, + skipped_aggregation_rows, + )) + } + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("PartialHashAggregateStream[{partition}]")) + .with_can_spill(true) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + batch_size, + baseline_metrics, + reservation, + reduction_factor, + skip_aggregation_probe, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + state: Some(PartialHashAggregateState::ReadingInput { hash_table }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> PartialHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialHashAggregateState::Error, + )) + } + + fn break_with_internal_err( + message: impl std::fmt::Display, + ) -> PartialHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit( + &self, + hash_table: &AggregateHashTable<PartialMarker>, + ) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= hash_table.building_group_count()) + } + + /// Updates skip aggregation probe state. + fn update_skip_aggregation_probe(&mut self, input_rows: usize, num_groups: usize) { + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.update_state(input_rows, num_groups); + } + } + + /// Returns true if the aggregation probe indicates that aggregation + /// should be skipped. + fn should_skip_aggregation(&self) -> bool { + self.skip_aggregation_probe + .as_ref() + .is_some_and(|probe| probe.should_skip()) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable<PartialMarker>, + close_input: bool, + ) -> Result<()> { + if close_input { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + hash_table.start_output() + } + + /// Handle ReadingInput state - aggregate input batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::ReadingInput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected ReadingInput state", + ); + }; + debug_assert!(hash_table.is_building()); + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + PartialHashAggregateState::ReadingInput { hash_table }, + )), + Poll::Ready(Some(Ok(batch))) => { + // ---------------------------------- + // Step 1: Aggregate the input batch + // ---------------------------------- + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + let result = hash_table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // -------------------------------- + // Step 2: Soft limit optimization + // -------------------------------- + if self.hit_soft_group_limit(&hash_table) { + let timer = elapsed_compute.timer(); + let result = self.start_output(&mut hash_table, true); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + return ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: None, + }, + ); + } + + // ---------------------------------------------- + // Step 3: Skip partial aggregation optimization + // ---------------------------------------------- + self.update_skip_aggregation_probe( + input_rows, + hash_table.building_group_count(), + ); + + // True branch: a decision has been made to skip partial aggregation. + if self.should_skip_aggregation() { + let timer = elapsed_compute.timer(); + let result = match hash_table.partial_skip_table() { + Ok(skip_hash_table) => self + .start_output(&mut hash_table, false) + .map(|()| skip_hash_table), + Err(e) => Err(e), + }; + timer.done(); + + match result { + Ok(skip_hash_table) => { + // Move to `ProducingOutput` first. Its `skip_hash_table` + // field moves the stream to skip-partial aggregation after + // the accumulated batches have been output. + return ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: Some(skip_hash_table), + }, + ); + } + Err(e) => return Self::break_with_err(e), + } + } + + // ------------------------------------------------- + // Step 4: Larger-than-memory execution (early emit) + // ------------------------------------------------- + let timer = elapsed_compute.timer(); + let resize_result = self.reservation.try_resize(hash_table.memory_size()); + timer.done(); + match resize_result { + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(_)) => { + let elapsed_compute = + self.baseline_metrics.elapsed_compute().clone(); + // Stops on drop + let _timer = elapsed_compute.timer(); + let state_batch_result = hash_table.take_state_batch(); + + // Emitting clears the aggregate table and releases its + // accumulated memory. Update the reservation accordingly. + let resize_result = + self.reservation.try_resize(hash_table.memory_size()); + + if let Err(e) = resize_result { + return Self::break_with_err(e); + } + + let materialized_group_states = match state_batch_result { + Ok(Some(batch)) => batch, + Ok(None) => { + return Self::break_with_err(internal_datafusion_err!( + "Partial hash aggregate ran out of memory with no aggregated groups" + )); + } + Err(e) => return Self::break_with_err(e), + }; + + return ControlFlow::Continue( + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: materialized_group_states, + }, + ); + } + Err(e) => return Self::break_with_err(e), + } + + ControlFlow::Continue(PartialHashAggregateState::ReadingInput { + hash_table, + }) + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = self.start_output(&mut hash_table, true); + timer.done(); + + match result { + Ok(()) => ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: None, + }, + ), + Err(e) => Self::break_with_err(e), + } + } + } + } + + /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state + /// batch in `batch_size`(from configuration) slices, then 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: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: batch, + } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected EmittingOnMemoryPressure state", + ); + }; + + let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + // Last batch to output, go back to `ReadingInput` + ( + batch, + PartialHashAggregateState::ReadingInput { hash_table }, + ) + } else { + // More batch 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, + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: remaining, + }, + ) + }; + + self.reduction_factor.add_part(output_batch.num_rows()); + 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 partial aggregate state batches. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::ProducingOutput { + mut hash_table, + skip_hash_table, + } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected ProducingOutput state", + ); + }; + debug_assert!(!hash_table.is_building()); + + 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 _ = self.reservation.try_resize(hash_table.memory_size()); + self.reduction_factor.add_part(batch.num_rows()); + debug_assert!(batch.num_rows() > 0); + let next_state = if hash_table.is_done() { + match skip_hash_table { + Some(hash_table) => { + PartialHashAggregateState::SkippingAggregation { hash_table } + } + None => PartialHashAggregateState::Done, + } + } else { + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table, + } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Ok(None) => { + let _ = self.reservation.try_resize(0); + // If the previous `Aggregating` stage decided to skip partial + // aggregation, go to the `SkippingAggregation` stage; otherwise finish. + let next_state = match skip_hash_table { + Some(hash_table) => { + PartialHashAggregateState::SkippingAggregation { hash_table } + } + None => PartialHashAggregateState::Done, + }; + ControlFlow::Continue(next_state) + } + Err(e) => Self::break_with_err(e), + } + } + + /// Handle SkippingAggregation state - convert raw input directly to partial states. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_skipping_aggregation( + &mut self, + cx: &mut Context<'_>, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::SkippingAggregation { mut hash_table } = + original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected SkippingAggregation state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + PartialHashAggregateState::SkippingAggregation { hash_table }, + )), + Poll::Ready(Some(Ok(batch))) => { + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.record_skipped(&batch); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.convert_batch_to_state(&batch); + timer.done(); + + match result { + Ok(batch) => ControlFlow::Break(( + Poll::Ready(Some( + Ok(batch.record_output(&self.baseline_metrics)), + )), + PartialHashAggregateState::SkippingAggregation { hash_table }, + )), + Err(e) => Self::break_with_err(e), + } + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + ControlFlow::Continue(PartialHashAggregateState::Done) + } + } + } +} + +impl Stream for PartialHashAggregateStream { + type Item = Result<RecordBatch>; + + /// Entry point for the partial hash aggregate state machine. + /// + /// See comments in [`PartialHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling input and aggregating batches into the + /// in-memory hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one 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(skip=None) + /// Input was exhausted, or the soft group limit was reached. Move to + /// the next state to start outputting. + /// -> ProducingOutput(skip=Some) + /// Partial skip aggregation was triggered. First move to the + /// `ProducingOutput` state to drain the accumulated state, then move to + /// the `SkippingAggregation` state to convert input directly to partial + /// state without aggregation. + /// + /// 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(skip=None) + /// -> ProducingOutput(skip=None) + /// One accumulated output batch was yielded, repeat to continue producing + /// output incrementally. + /// -> Done + /// All accumulated output was emitted. + /// + /// ProducingOutput(skip=Some) + /// -> ProducingOutput(skip=Some) + /// One accumulated output batch was yielded, repeat to continue producing + /// output incrementally. + /// -> SkippingAggregation + /// All accumulated output was emitted. Continue by converting raw + /// input batches directly to partial aggregate state. + /// + /// SkippingAggregation + /// -> SkippingAggregation + /// One `convert_to_state` batch was yielded; repeat to continue + /// processing. + /// -> Done + /// Input was exhausted. + /// + /// 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<Option<Self::Item>> { + loop { + let cur_state = self + .state + .take() + .expect("PartialHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ PartialHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ PartialHashAggregateState::EmittingOnMemoryPressure { .. } => { + self.handle_emitting_on_memory_pressure(state) + } + state @ PartialHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ PartialHashAggregateState::SkippingAggregation { .. } => { + self.handle_skipping_aggregation(cx, state) + } + state @ PartialHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } + state @ PartialHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, PartialHashAggregateState::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(PartialHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for PartialHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl FinalHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug_assert!(matches!( + agg.mode, + super::AggregateMode::Final | super::AggregateMode::FinalPartitioned + )); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + + let hash_table = AggregateHashTable::<FinalMarker>::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + + let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + Some(Box::new(FinalSpillContext::new( + agg, + context, + partition, + batch_size, + &input_schema, + spill_metrics, + )?)) + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("FinalHashAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + state: Some(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> FinalHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + FinalHashAggregateState::Error, + )) + } + + fn break_with_internal_err( + message: impl std::fmt::Display, + ) -> FinalHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable<FinalMarker>) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= hash_table.building_group_count()) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable<FinalMarker>, + ) -> Result<()> { + self.close_input(); + hash_table.start_output() + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + hash_table: &AggregateHashTable<FinalMarker>, + spill_context: Option<&FinalSpillContext>, + ) -> usize { + let table_size = hash_table.memory_size(); + if spill_context.is_some() { + // Count extra space needed for in-memory sorting and spilling. Only + // count memory for indices, the payload will be materialize incrementally + // in smaller chunks. + table_size.saturating_add( + hash_table + .building_group_count() + .saturating_mul(size_of::<u32>()), + ) + } else { + table_size + } + } + + /// Handle ReadingInput state - aggregate partial state batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_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", + ); + }; + + 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(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // 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(); + + return match result { + Ok(()) => ControlFlow::Continue( + FinalHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + }; + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + 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), + } + + 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), + } + } + } + } + } + } + + /// 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( + &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", + ); + } + + 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")); + } + + 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), + } + } + + /// 1. Spills the last in-memory run. + /// 2. 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 + /// 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( + &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 group_by_metrics = hash_table.group_by_metrics().clone(); + drop(hash_table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_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), + } + } + + /// 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( + &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", + ); + }; + + 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<RecordBatch>; + + /// 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<Option<Self::Item>> { + 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); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, FinalHashAggregateState::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(FinalHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for FinalHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::aggregates::{AggregateMode, PhysicalGroupBy}; + use crate::execution_plan::ExecutionPlan; + use crate::test::TestMemoryExec; + + use arrow::array::{Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use futures::StreamExt; + + #[tokio::test] + async fn test_partial_hash_stream_double_emission_race_condition_bug() -> Result<()> { + // Fix for https://github.com/apache/datafusion/issues/18701 + // This test specifically proves that we have fixed double emission race condition + // where emit_early_if_necessary() and switch_to_skip_aggregation() + // both emit in the same loop iteration, causing data loss + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // Create data that will trigger BOTH conditions in the same iteration: + // 1. More groups than batch_size (triggers early emission when memory pressure hits) + // 2. High cardinality ratio (triggers skip aggregation) + let batch_size = 1024; // We'll set this in session config + let num_groups = batch_size + 100; // Slightly more than batch_size (1124 groups) + + // Create exactly 1 row per group = 100% cardinality ratio + let group_ids: Vec<i32> = (0..num_groups as i32).collect(); + let values: Vec<i64> = vec![1; num_groups]; + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids)), + Arc::new(Int64Array::from(values)), + ], + )?; + let input_partitions = vec![vec![batch]]; + + // Create constrained memory to trigger early emission but not completely fail + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1024, 1.0) // small enough to start but will trigger pressure + .build_arc()?; + + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure to trigger BOTH conditions: + // 1. Low probe threshold (triggers skip probe after few rows) + // 2. Low ratio threshold (triggers skip aggregation immediately) + // 3. Set batch_size to 1024 so our 1124 groups will trigger early emission + // This creates the race condition where both emit paths are triggered + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.batch_size", + &datafusion_common::ScalarValue::UInt64(Some(1024)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(50)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(0.8)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode where the race condition occurs + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Count total groups emitted + let mut total_output_groups = 0; + for batch in &results { + total_output_groups += batch.num_rows(); + } + + assert_eq!( + total_output_groups, num_groups, + "Unexpected number of groups", + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partial_hash_stream_skip_aggregation_probe_not_locked_until_skip() + -> Result<()> { + // Test that the probe is not locked until we actually decide to skip. + // This allows us to continue evaluating the skip condition across multiple batches. + // + // Scenario: + // - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip + // - Batch 2: Now hits ratio threshold (high cardinality) -> skip + // + // Without the fix, the probe would be locked after batch 1, preventing the skip + // decision from being made on batch 2. + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int32, false), + ])); + + // Configure thresholds: + // - probe_rows_threshold: 100 rows + // - probe_ratio_threshold: 0.8 (80%) + let probe_rows_threshold = 100; + let probe_ratio_threshold = 0.8; + + // Batch 1: 100 rows with only 10 unique groups + // Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip + // This will hit the rows threshold but not the ratio threshold + let batch1_rows = 100; + let batch1_groups = 10; + let mut group_ids_batch1 = Vec::new(); + for i in 0..batch1_rows { + group_ids_batch1.push((i % batch1_groups) as i32); + } + let values_batch1: Vec<i32> = vec![1; batch1_rows]; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch1)), + Arc::new(Int32Array::from(values_batch1)), + ], + )?; + + // Batch 2: 360 rows with 360 unique NEW groups (starting from group 10) + // After batch 2, total: 460 rows, 370 groups + // Ratio: 370/460 is about 0.804 (80.4%) > 0.8 -> SHOULD decide to skip + let batch2_rows = 360; + let batch2_groups = 360; + let group_ids_batch2: Vec<i32> = (batch1_groups..(batch1_groups + batch2_groups)) + .map(|x| x as i32) + .collect(); + let values_batch2: Vec<i32> = vec![1; batch2_rows]; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch2)), + Arc::new(Int32Array::from(values_batch2)), + ], + )?; + + // Batch 3: This batch should be skipped since we decided to skip after batch 2 + // 100 rows with 100 unique groups (continuing from where batch 2 left off) + let batch3_rows = 100; + let batch3_groups = 100; + let batch3_start_group = batch1_groups + batch2_groups; + let group_ids_batch3: Vec<i32> = (batch3_start_group + ..(batch3_start_group + batch3_groups)) + .map(|x| x as i32) + .collect(); + let values_batch3: Vec<i32> = vec![1; batch3_rows]; + + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch3)), + Arc::new(Int32Array::from(values_batch3)), + ], + )?; + + let input_partitions = vec![vec![batch1, batch2, batch3]]; + + let runtime = RuntimeEnvBuilder::default().build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure skip aggregation settings + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Check that skip aggregation actually happened. + // The key metric is skipped_aggregation_rows. + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + // We expect batch 3's rows to be skipped (100 rows) + assert_eq!( + skipped_rows, batch3_rows, + "Expected batch 3's rows ({batch3_rows}) to be skipped", + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs new file mode 100644 index 0000000..a39c6f3 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/mod.rs
@@ -0,0 +1,7982 @@ +// 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. + +//! Aggregate functionality +//! +//! # Aggregate planning +//! +//! DataFusion selects different aggregate implementations (streams) based on the +//! query shape and configuration. This section provides an overview of the +//! available stream variants. +//! +//! See each stream's documentation for details. +//! +//! ## 1. Two-stage hash aggregation +//! +//! Two-stage hash aggregation is used for regular parallel execution. +//! +//! The input passes through three execution operators to produce the final +//! aggregation result: +//! +//! 1. Partial aggregation reads the input and produces partial states. It +//! aggregates independently within each partition, which usually reduces +//! cardinality before the later shuffle. +//! 2. Hash repartitioning on the group keys sends all partial states for each +//! group to the same output partition for final aggregation. +//! 3. Final aggregation reads the partial states, combines them, and emits the +//! final results. +//! +//! ```text +//! AggregateExec (final) +//! RepartitionExec (hash by group keys) +//! AggregateExec (partial) +//! ``` +//! +//! See [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] for details. +//! +//! ### Ordering optimization +//! +//! When the input is ordered by the group key, an ordered fast path is used. It +//! uses a similar two-stage hash aggregation with an early-emission optimization. +//! +//! ```text +//! AggregateExec (final, ordered) +//! RepartitionExec (hash by group keys, order-preserving) +//! AggregateExec (partial, ordered) +//! ``` +//! +//! See [`OrderedPartialAggregateStream`] and [`OrderedFinalAggregateStream`] for +//! details. +//! +//! Related configuration: +//! +//! - [`datafusion.execution.target_partitions`](datafusion_common::config::ExecutionOptions::target_partitions) +//! - [`datafusion.optimizer.repartition_aggregations`](datafusion_common::config::OptimizerOptions::repartition_aggregations) +//! - [`datafusion.optimizer.prefer_existing_sort`](datafusion_common::config::OptimizerOptions::prefer_existing_sort) +//! +//! ## 2. Single-stage hash aggregation +//! +//! When there is a single partition, or the aggregation input is already +//! key-partitioned (e.g., a data source has existing range partitioning), +//! `Single` mode aggregation is used. +//! +//! It takes raw input and directly produces the final result. +//! +//! ```text +//! AggregateExec (mode=Single or SinglePartitioned) +//! input +//! ``` +//! +//! See [`SingleHashAggregateStream`] for details. +//! +//! Related configuration: +//! +//! - [`datafusion.execution.target_partitions`](datafusion_common::config::ExecutionOptions::target_partitions) +//! - [`datafusion.optimizer.repartition_aggregations`](datafusion_common::config::OptimizerOptions::repartition_aggregations) +//! +//! ## 3. Aggregation without grouping expressions +//! +//! A global aggregate maintains one accumulator set per input partition rather +//! than a hash table of groups. Partial stages compute local states and a final +//! stage combines them into one output row: +//! +//! ```text +//! AggregateExec (final, no-grouping) +//! CoalescePartitionsExec +//! AggregateExec (partial, no-grouping) +//! ``` +//! +//! Every stage without grouping expressions uses [`AggregateStream`]. This path +//! is selected before the grouped-stream migration setting is considered. +//! +//! ## 4. Grouped TopK aggregation +//! +//! When a query only needs the best `N` groups, retaining every group in a hash +//! table and sorting them afterward does unnecessary work. The optimizer pushes +//! the sort limit and direction into the aggregate: +//! +//! ```text +//! SortExec (fetch=N) +//! AggregateExec (limit=N, order=...) +//! input +//! ``` +//! +//! [`GroupedTopKAggregateStream`] keeps a bounded priority map for a single group +//! key. It supports group-by-only queries and compatible `MIN` or `MAX` +//! aggregates. An unordered group-by-only soft limit instead stays on the normal +//! hash aggregation path. +//! +//! Related configuration: +//! +//! - [`datafusion.optimizer.enable_topk_aggregation`](datafusion_common::config::OptimizerOptions::enable_topk_aggregation) +//! - [`datafusion.optimizer.enable_distinct_aggregation_soft_limit`](datafusion_common::config::OptimizerOptions::enable_distinct_aggregation_soft_limit) +//! +//! ## 5. Partial-reduce hash aggregation +//! +//! This implementation will not be planned by DataFusion SQL interface, it must be +//! manually constructed at [`ExecutionPlan`] level. +//! +//! This mode is useful in a distributed setting. +//! +//! See [`PartialReduceHashAggregateStream`] for details. +//! +//! ## 6. Fallback grouped hash aggregation +//! +//! [`GroupedHashAggregateStream`] is the legacy implementation for several of the +//! stream types above. It is being incrementally migrated to separate streams. +//! +//! See the issue for details: <https://github.com/apache/datafusion/issues/22710> +#![expect(rustdoc::private_intra_doc_links)] + +use std::borrow::Cow; +use std::sync::Arc; + +use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; +use crate::aggregates::{ + aggregate_stream::AggregateStream, + grouped_hash_stream::GroupedHashAggregateStream, + grouped_topk_stream::GroupedTopKAggregateStream, + hash_stream::{FinalHashAggregateStream, PartialHashAggregateStream}, + ordered_final_stream::OrderedFinalAggregateStream, + ordered_partial_stream::OrderedPartialAggregateStream, + partial_reduce_stream::PartialReduceHashAggregateStream, + single_stream::SingleHashAggregateStream, +}; +use crate::execution_plan::{ + CardinalityEffect, EmissionType, plan_contains_expression_id, +}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; +use crate::{ + DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + InputOrderMode, SendableRecordBatchStream, Statistics, +}; +use datafusion_common::config::ConfigOptions; +use parking_lot::Mutex; +use std::collections::{HashMap, HashSet}; + +use arrow::array::{ArrayRef, UInt8Array, UInt16Array, UInt32Array, UInt64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use arrow_schema::FieldRef; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + ColumnStatistics, Constraint, Constraints, Result, ScalarValue, + 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; +use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::{ + ConstExpr, EquivalenceProperties, physical_exprs_contains, +}; +use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, fmt_sql}; +use datafusion_physical_expr_common::sort_expr::{ + LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortRequirement, +}; + +use datafusion_expr::utils::AggregateOrderSensitivity; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use itertools::Itertools; +use topk::hash_table::is_supported_hash_key_type; +use topk::heap::is_supported_heap_type; + +mod aggregate_hash_table; +mod aggregate_stream; +pub mod group_values; +mod grouped_hash_stream; +mod grouped_topk_stream; +mod hash_stream; +pub mod order; +mod ordered_final_stream; +mod ordered_partial_stream; +mod partial_reduce_stream; +mod single_stream; +mod skip_partial; +mod topk; + +/// Returns true if TopK aggregation data structures support the provided key and value types. +/// +/// This function checks whether both the key type (used for grouping) and value type +/// (used in min/max aggregation) can be handled by the TopK aggregation heap and hash table. +/// Supported types include Arrow primitives (integers, floats, decimals, intervals) and +/// UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`). +/// ```text +pub fn topk_types_supported(key_type: &DataType, value_type: &DataType) -> bool { + is_supported_hash_key_type(key_type) && is_supported_heap_type(value_type) +} + +/// Hard-coded seed for aggregations to ensure hash values differ from `RepartitionExec`, avoiding collisions. +const AGGREGATION_HASH_SEED: datafusion_common::hash_utils::RandomState = + // This seed is chosen to be a large 64-bit number + datafusion_common::hash_utils::RandomState::with_seed(15395726432021054657); + +/// Whether an aggregate stage consumes raw input data or intermediate +/// accumulator state from a previous aggregation stage. +/// +/// See the [table on `AggregateMode`](AggregateMode#variants-and-their-inputoutput-modes) +/// for how this relates to aggregate modes. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum AggregateInputMode { + /// The stage consumes raw, unaggregated input data and calls + /// [`Accumulator::update_batch`]. + Raw, + /// The stage consumes intermediate accumulator state from a previous + /// aggregation stage and calls [`Accumulator::merge_batch`]. + Partial, +} + +/// Whether an aggregate stage produces intermediate accumulator state +/// or final output values. +/// +/// See the [table on `AggregateMode`](AggregateMode#variants-and-their-inputoutput-modes) +/// for how this relates to aggregate modes. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum AggregateOutputMode { + /// The stage produces intermediate accumulator state, serialized via + /// [`Accumulator::state`]. + Partial, + /// The stage produces final output values via + /// [`Accumulator::evaluate`]. + Final, +} + +/// Aggregation modes +/// +/// See [`Accumulator::state`] for background information on multi-phase +/// aggregation and how these modes are used. +/// +/// # Variants and their input/output modes +/// +/// Each variant can be characterized by its [`AggregateInputMode`] and +/// [`AggregateOutputMode`]: +/// +/// ```text +/// | Input: Raw data | Input: Partial state +/// Output: Final values | Single, SinglePartitioned | Final, FinalPartitioned +/// Output: Partial state | Partial | PartialReduce +/// ``` +/// +/// Use [`AggregateMode::input_mode`] and [`AggregateMode::output_mode`] +/// to query these properties. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum AggregateMode { + /// One of multiple layers of aggregation, any input partitioning + /// + /// Partial aggregate that can be applied in parallel across input + /// partitions. + /// + /// This is the first phase of a multi-phase aggregation. + Partial, + /// *Final* of multiple layers of aggregation, in exactly one partition + /// + /// Final aggregate that produces a single partition of output by combining + /// the output of multiple partial aggregates. + /// + /// This is the second phase of a multi-phase aggregation. + /// + /// This mode requires that the input is a single partition + /// + /// Note: Adjacent `Partial` and `Final` mode aggregation is equivalent to a `Single` + /// mode aggregation node. The `Final` mode is required since this is used in an + /// intermediate step. The [`CombinePartialFinalAggregate`] physical optimizer rule + /// will replace this combination with `Single` mode for more efficient execution. + /// + /// [`CombinePartialFinalAggregate`]: https://docs.rs/datafusion/latest/datafusion/physical_optimizer/combine_partial_final_agg/struct.CombinePartialFinalAggregate.html + Final, + /// *Final* of multiple layers of aggregation, input is *Partitioned* + /// + /// Final aggregate that works on pre-partitioned data. + /// + /// This mode requires that all rows with a particular grouping key are in + /// the same partitions, such as is the case with Hash repartitioning on the + /// group keys. If a group key is duplicated, duplicate groups would be + /// produced + FinalPartitioned, + /// *Single* layer of Aggregation, input is exactly one partition + /// + /// Applies the entire logical aggregation operation in a single operator, + /// as opposed to Partial / Final modes which apply the logical aggregation using + /// two operators. + /// + /// This mode requires that the input is a single partition (like Final) + Single, + /// *Single* layer of Aggregation, input is *Partitioned* + /// + /// Applies the entire logical aggregation operation in a single operator, + /// as opposed to Partial / Final modes which apply the logical aggregation + /// using two operators. + /// + /// This mode requires that the input has more than one partition, and is + /// partitioned by group key (like FinalPartitioned). + SinglePartitioned, + /// Combine multiple partial aggregations to produce a new partial + /// aggregation. + /// + /// Input is intermediate accumulator state (like Final), but output is + /// also intermediate accumulator state (like Partial). This enables + /// tree-reduce aggregation strategies where partial results from + /// multiple workers are combined in multiple stages before a final + /// evaluation. + /// + /// ```text + /// Final + /// / \ + /// PartialReduce PartialReduce + /// / \ / \ + /// Partial Partial Partial Partial + /// ``` + /// + /// # Motivation + /// + /// This reduces shuffling traffic in a distributed setting. See + /// <https://github.com/datafusion-contrib/datafusion-distributed/issues/360> + /// for details. + PartialReduce, +} + +impl AggregateMode { + /// Returns the [`AggregateInputMode`] for this mode: whether this + /// stage consumes raw input data or intermediate accumulator state. + /// + /// See the [table above](AggregateMode#variants-and-their-inputoutput-modes) + /// for details. + pub fn input_mode(&self) -> AggregateInputMode { + match self { + AggregateMode::Partial + | AggregateMode::Single + | AggregateMode::SinglePartitioned => AggregateInputMode::Raw, + AggregateMode::Final + | AggregateMode::FinalPartitioned + | AggregateMode::PartialReduce => AggregateInputMode::Partial, + } + } + + /// Returns the [`AggregateOutputMode`] for this mode: whether this + /// stage produces intermediate accumulator state or final output values. + /// + /// See the [table above](AggregateMode#variants-and-their-inputoutput-modes) + /// for details. + pub fn output_mode(&self) -> AggregateOutputMode { + match self { + AggregateMode::Final + | AggregateMode::FinalPartitioned + | AggregateMode::Single + | AggregateMode::SinglePartitioned => AggregateOutputMode::Final, + AggregateMode::Partial | AggregateMode::PartialReduce => { + AggregateOutputMode::Partial + } + } + } +} + +/// Represents `GROUP BY` clause in the plan (including the more general GROUPING SET) +/// In the case of a simple `GROUP BY a, b` clause, this will contain the expression [a, b] +/// and a single group [false, false]. +/// In the case of `GROUP BY GROUPING SETS/CUBE/ROLLUP` the planner will expand the expression +/// into multiple groups, using null expressions to align each group. +/// For example, with a group by clause `GROUP BY GROUPING SETS ((a,b),(a),(b))` the planner should +/// create a `PhysicalGroupBy` like +/// ```text +/// PhysicalGroupBy { +/// expr: [(col(a), a), (col(b), b)], +/// null_expr: [(NULL, a), (NULL, b)], +/// groups: [ +/// [false, false], // (a,b) +/// [false, true], // (a) <=> (a, NULL) +/// [true, false] // (b) <=> (NULL, b) +/// ] +/// } +/// ``` +#[derive(Clone, Debug, Default)] +pub struct PhysicalGroupBy { + /// Distinct (Physical Expr, Alias) in the grouping set + expr: Vec<(Arc<dyn PhysicalExpr>, String)>, + /// Corresponding NULL expressions for expr + null_expr: Vec<(Arc<dyn PhysicalExpr>, String)>, + /// Null mask for each group in this grouping set. Each group is + /// composed of either one of the group expressions in expr or a null + /// expression in null_expr. If `groups[i][j]` is true, then the + /// j-th expression in the i-th group is NULL, otherwise it is `expr[j]`. + groups: Vec<Vec<bool>>, + /// True when GROUPING SETS/CUBE/ROLLUP are used so `__grouping_id` should + /// be included in the output schema. + has_grouping_set: bool, +} + +impl PhysicalGroupBy { + /// Create a new `PhysicalGroupBy` + pub fn new( + expr: Vec<(Arc<dyn PhysicalExpr>, String)>, + null_expr: Vec<(Arc<dyn PhysicalExpr>, String)>, + groups: Vec<Vec<bool>>, + has_grouping_set: bool, + ) -> Self { + Self { + expr, + null_expr, + groups, + has_grouping_set, + } + } + + /// Create a GROUPING SET with only a single group. This is the "standard" + /// case when building a plan from an expression such as `GROUP BY a,b,c` + pub fn new_single(expr: Vec<(Arc<dyn PhysicalExpr>, String)>) -> Self { + let num_exprs = expr.len(); + Self { + expr, + null_expr: vec![], + groups: vec![vec![false; num_exprs]], + has_grouping_set: false, + } + } + + /// Calculate GROUP BY expressions nullable + pub fn exprs_nullable(&self) -> Vec<bool> { + let mut exprs_nullable = vec![false; self.expr.len()]; + for group in self.groups.iter() { + group.iter().enumerate().for_each(|(index, is_null)| { + if *is_null { + exprs_nullable[index] = true; + } + }) + } + exprs_nullable + } + + /// Returns true if this has no grouping at all (including no GROUPING SETS) + pub fn is_true_no_grouping(&self) -> bool { + self.is_empty() && !self.has_grouping_set + } + + /// Returns the group expressions + pub fn expr(&self) -> &[(Arc<dyn PhysicalExpr>, String)] { + &self.expr + } + + /// Returns the null expressions + pub fn null_expr(&self) -> &[(Arc<dyn PhysicalExpr>, String)] { + &self.null_expr + } + + /// Returns the group null masks + pub fn groups(&self) -> &[Vec<bool>] { + &self.groups + } + + /// Returns true if this grouping uses GROUPING SETS, CUBE or ROLLUP. + pub fn has_grouping_set(&self) -> bool { + self.has_grouping_set + } + + /// Returns true if this `PhysicalGroupBy` has no group expressions + pub fn is_empty(&self) -> bool { + self.expr.is_empty() + } + + /// Returns true if this is a "simple" GROUP BY (not using GROUPING SETS/CUBE/ROLLUP). + /// This determines whether the `__grouping_id` column is included in the output schema. + pub fn is_single(&self) -> bool { + !self.has_grouping_set + } + + /// Calculate GROUP BY expressions according to input schema. + pub fn input_exprs(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.expr + .iter() + .map(|(expr, _alias)| Arc::clone(expr)) + .collect() + } + + /// The number of expressions in the output schema. + fn num_output_exprs(&self) -> usize { + let mut num_exprs = self.expr.len(); + if self.has_grouping_set { + num_exprs += 1 + } + num_exprs + } + + /// Return grouping expressions as they occur in the output schema. + pub fn output_exprs(&self) -> Vec<Arc<dyn PhysicalExpr>> { + let num_output_exprs = self.num_output_exprs(); + let mut output_exprs = Vec::with_capacity(num_output_exprs); + output_exprs.extend( + self.expr + .iter() + .enumerate() + .take(num_output_exprs) + .map(|(index, (_, name))| Arc::new(Column::new(name, index)) as _), + ); + if self.has_grouping_set { + output_exprs.push(Arc::new(Column::new( + Aggregate::INTERNAL_GROUPING_ID, + self.expr.len(), + )) as _); + } + output_exprs + } + + /// Returns the number expression as grouping keys. + pub fn num_group_exprs(&self) -> usize { + self.expr.len() + usize::from(self.has_grouping_set) + } + + /// Returns the Arrow data type of the `__grouping_id` column. + /// + /// The type is chosen to be wide enough to hold both the semantic bitmask + /// (in the low `n` bits, where `n` is the number of grouping expressions) + /// and the duplicate ordinal (in the high bits). + fn grouping_id_data_type(&self) -> DataType { + Aggregate::grouping_id_type(self.expr.len(), max_duplicate_ordinal(&self.groups)) + } + + pub fn group_schema(&self, schema: &Schema) -> Result<SchemaRef> { + Ok(Arc::new(Schema::new(self.group_fields(schema)?))) + } + + /// Returns the fields that are used as the grouping keys. + fn group_fields(&self, input_schema: &Schema) -> Result<Vec<FieldRef>> { + let mut fields = Vec::with_capacity(self.num_group_exprs()); + for ((expr, name), group_expr_nullable) in + self.expr.iter().zip(self.exprs_nullable()) + { + fields.push( + Field::new( + name, + expr.data_type(input_schema)?, + group_expr_nullable || expr.nullable(input_schema)?, + ) + .with_metadata(expr.return_field(input_schema)?.metadata().clone()) + .into(), + ); + } + if self.has_grouping_set { + fields.push( + Field::new( + Aggregate::INTERNAL_GROUPING_ID, + self.grouping_id_data_type(), + false, + ) + .into(), + ); + } + Ok(fields) + } + + /// Returns the output fields of the group by. + /// + /// This might be different from the `group_fields` that might contain internal expressions that + /// should not be part of the output schema. + fn output_fields(&self, input_schema: &Schema) -> Result<Vec<FieldRef>> { + let mut fields = self.group_fields(input_schema)?; + fields.truncate(self.num_output_exprs()); + Ok(fields) + } + + /// Returns the `PhysicalGroupBy` for a final aggregation if `self` is used for a partial + /// aggregation. + pub fn as_final(&self) -> PhysicalGroupBy { + let expr: Vec<_> = + self.output_exprs() + .into_iter() + .zip( + self.expr.iter().map(|t| t.1.clone()).chain(std::iter::once( + Aggregate::INTERNAL_GROUPING_ID.to_owned(), + )), + ) + .collect(); + let num_exprs = expr.len(); + let groups = if self.expr.is_empty() && !self.has_grouping_set { + // No GROUP BY expressions - should have no groups + vec![] + } else { + vec![vec![false; num_exprs]] + }; + Self { + expr, + null_expr: vec![], + groups, + has_grouping_set: false, + } + } +} + +impl PartialEq for PhysicalGroupBy { + fn eq(&self, other: &PhysicalGroupBy) -> bool { + self.expr.len() == other.expr.len() + && self + .expr + .iter() + .zip(other.expr.iter()) + .all(|((expr1, name1), (expr2, name2))| expr1.eq(expr2) && name1 == name2) + && self.null_expr.len() == other.null_expr.len() + && self + .null_expr + .iter() + .zip(other.null_expr.iter()) + .all(|((expr1, name1), (expr2, name2))| expr1.eq(expr2) && name1 == name2) + && self.groups == other.groups + && self.has_grouping_set == other.has_grouping_set + } +} + +/// Streams used by [`AggregateExec`]. +/// +/// # Stream Variant Schema Notation +/// For example, `SELECT g, AVG(x) FROM t GROUP BY g` uses these schemas: +/// +/// ```text +/// initial input: [g, x] +/// partial state: [g, AVG(x) state columns, e.g. sum/count] +/// final result: [g, AVG(x)] +/// ``` +#[expect(clippy::large_enum_variant)] +enum StreamType { + /// Single group (no group by) aggregate stream. + /// Input output scheme: initial input -> final result + AggregateStream(AggregateStream), + /// Partial stage of the hash aggregation + /// Input output scheme: initial input -> partial state + PartialHash(PartialHashAggregateStream), + /// Partial-reduce stage of the hash aggregation + /// Input output scheme: partial state -> partial state + PartialReduceHash(PartialReduceHashAggregateStream), + /// Final stage of the hash aggregation + /// Input output scheme: partial state -> final result + FinalHash(FinalHashAggregateStream), + /// Single stage of the hash aggregation + /// Input output scheme: initial input -> final result + SingleHash(SingleHashAggregateStream), + /// Partial stage of aggregation for ordered input. + OrderedPartialAggregate(OrderedPartialAggregateStream), + /// Final stage of aggregation for ordered input. + OrderedFinalAggregate(OrderedFinalAggregateStream), + /// Hash aggregation reused for multiple stages + /// + /// Note this is being incrementally migrated to dedicated streams like + /// [`StreamType::PartialHash`], [`StreamType::FinalHash`], + /// [`StreamType::OrderedPartialAggregate`], and + /// [`StreamType::OrderedFinalAggregate`] + /// + /// See issue for details: <https://github.com/apache/datafusion/issues/22710> + GroupedHash(GroupedHashAggregateStream), + /// Grouped TopK aggregate stream. + /// Input output scheme: initial input -> final result + /// + /// Used for grouped aggregation with LIMIT / ordering, where the stream keeps + /// only the top groups required by the query. + GroupedPriorityQueue(GroupedTopKAggregateStream), +} + +impl From<StreamType> for SendableRecordBatchStream { + fn from(stream: StreamType) -> Self { + match stream { + 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::SingleHash(stream) => Box::pin(stream), + StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), + StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), + StreamType::GroupedHash(stream) => Box::pin(stream), + StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), + } + } +} + +/// # Aggregate Dynamic Filter Pushdown Overview +/// +/// For queries like +/// -- `example_table(type TEXT, val INT)` +/// SELECT min(val) +/// FROM example_table +/// WHERE type='A'; +/// +/// And `example_table`'s physical representation is a partitioned parquet file with +/// column statistics +/// - part-0.parquet: val {min=0, max=100} +/// - part-1.parquet: val {min=100, max=200} +/// - ... +/// - part-100.parquet: val {min=10000, max=10100} +/// +/// After scanning the 1st file, we know we only have to read files if their minimal +/// value on `val` column is less than 0, the minimal `val` value in the 1st file. +/// +/// We can skip scanning the remaining file by implementing dynamic filter, the +/// intuition is we keep a shared data structure for current min in both `AggregateExec` +/// and `DataSourceExec`, and let it update during execution, so the scanner can +/// know during execution if it's possible to skip scanning certain files. See +/// physical optimizer rule `FilterPushdown` for details. +/// +/// # Implementation +/// +/// ## Enable Condition +/// - No grouping (no `GROUP BY` clause in the sql, only a single global group to aggregate) +/// - The aggregate expression must be `min`/`max`, and evaluate directly on columns. +/// Note multiple aggregate expressions that satisfy this requirement are allowed, +/// and a dynamic filter will be constructed combining all applicable expr's +/// states. See more in the following example with dynamic filter on multiple columns. +/// +/// ## Filter Construction +/// The filter is kept in the `DataSourceExec`, and it will gets update during execution, +/// the reader will interpret it as "the upstream only needs rows that such filter +/// predicate is evaluated to true", and certain scanner implementation like `parquet` +/// can evaluate column statistics on those dynamic filters, to decide if they can +/// prune a whole range. +/// +/// ### Examples +/// - Expr: `min(a)`, Dynamic Filter: `a < a_cur_min` +/// - Expr: `min(a), max(a), min(b)`, Dynamic Filter: `(a < a_cur_min) OR (a > a_cur_max) OR (b < b_cur_min)` +#[derive(Debug, Clone)] +struct AggrDynFilter { + /// The physical expr for the dynamic filter shared between the `AggregateExec` + /// and the parquet scanner. + filter: Arc<DynamicFilterPhysicalExpr>, + /// The current bounds for the dynamic filter, updates during the execution to + /// tighten the bound for more effective pruning. + /// + /// Each vector element is for the accumulators that support dynamic filter. + /// e.g. This `AggregateExec` has accumulator: + /// min(a), avg(a), max(b) + /// And this field stores [PerAccumulatorDynFilter(min(a)), PerAccumulatorDynFilter(min(b))] + supported_accumulators_info: Vec<PerAccumulatorDynFilter>, +} + +// ---- Aggregate Dynamic Filter Utility Structs ---- + +/// Aggregate expressions that support the dynamic filter pushdown in aggregation. +/// See comments in [`AggrDynFilter`] for conditions. +#[derive(Debug, Clone)] +struct PerAccumulatorDynFilter { + aggr_type: DynamicFilterAggregateType, + /// During planning and optimization, the parent structure is kept in `AggregateExec`, + /// this index is into `aggr_expr` vec inside `AggregateExec`. + /// During execution, the parent struct is moved into `AggregateStream` (stream + /// for no grouping aggregate execution), and this index is into `aggregate_expressions` + /// vec inside `AggregateStreamInner` + aggr_index: usize, + // The current bound. Shared among all streams. + shared_bound: Arc<Mutex<ScalarValue>>, +} + +/// Aggregate types that are supported for dynamic filter in `AggregateExec` +#[derive(Debug, Clone)] +enum DynamicFilterAggregateType { + Min, + Max, +} + +/// Configuration for limit-based optimizations in aggregation +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct LimitOptions { + /// The maximum number of rows to return + pub limit: usize, + /// Optional ordering direction (true = descending, false = ascending) + /// This is used for TopK aggregation to maintain a priority queue with the correct ordering + pub descending: Option<bool>, +} + +impl LimitOptions { + /// Create a new LimitOptions with a limit and no specific ordering + pub fn new(limit: usize) -> Self { + Self { + limit, + descending: None, + } + } + + /// Create a new LimitOptions with a limit and ordering direction + pub fn new_with_order(limit: usize, descending: bool) -> Self { + Self { + limit, + descending: Some(descending), + } + } + + pub fn limit(&self) -> usize { + self.limit + } + + pub fn descending(&self) -> Option<bool> { + self.descending + } +} + +/// Hash aggregate execution plan +#[derive(Debug, Clone)] +pub struct AggregateExec { + /// Aggregation mode (full, partial) + mode: AggregateMode, + /// Group by expressions + /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance. + group_by: Arc<PhysicalGroupBy>, + /// Aggregate expressions + /// The same reason to [`Arc`] it as for [`Self::group_by`]. + aggr_expr: Arc<[Arc<AggregateFunctionExpr>]>, + /// FILTER (WHERE clause) expression for each aggregate expression + /// The same reason to [`Arc`] it as for [`Self::group_by`]. + filter_expr: Arc<[Option<Arc<dyn PhysicalExpr>>]>, + /// Configuration for limit-based optimizations + limit_options: Option<LimitOptions>, + /// Input plan, could be a partial aggregate or the input to the aggregate + pub input: Arc<dyn ExecutionPlan>, + /// Schema after the aggregate is applied. Contains the group by columns followed by the + /// aggregate outputs. + schema: SchemaRef, + /// Input schema before any aggregation is applied. For partial aggregate this will be the + /// same as input.schema() but for the final aggregate it will be the same as the input + /// to the partial aggregate, i.e., partial and final aggregates have same `input_schema`. + /// We need the input schema of partial aggregate to be able to deserialize aggregate + /// expressions from protobuf for final aggregate. + pub input_schema: SchemaRef, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + required_input_ordering: Option<OrderingRequirements>, + /// Describes how the input is ordered relative to the group by columns + input_order_mode: InputOrderMode, + cache: Arc<PlanProperties>, + /// During initialization, if the plan supports dynamic filtering (see [`AggrDynFilter`]), + /// it is set to `Some(..)` regardless of whether it can be pushed down to a child node. + /// + /// During filter pushdown optimization, if a child node can accept this filter, + /// it remains `Some(..)` to enable dynamic filtering during aggregate execution; + /// otherwise, it is cleared to `None`. + dynamic_filter: Option<Arc<AggrDynFilter>>, +} + +impl AggregateExec { + /// Function used in `OptimizeAggregateOrder` optimizer rule, + /// where we need parts of the new value, others cloned from the old one + /// Rewrites aggregate exec with new aggregate expressions. + pub fn with_new_aggr_exprs( + &self, + aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>, + ) -> Self { + Self { + aggr_expr: aggr_expr.into(), + // clone the rest of the fields + required_input_ordering: self.required_input_ordering.clone(), + metrics: ExecutionPlanMetricsSet::new(), + input_order_mode: self.input_order_mode.clone(), + cache: Arc::clone(&self.cache), + mode: self.mode, + group_by: Arc::clone(&self.group_by), + filter_expr: Arc::clone(&self.filter_expr), + limit_options: self.limit_options, + input: Arc::clone(&self.input), + schema: Arc::clone(&self.schema), + input_schema: Arc::clone(&self.input_schema), + dynamic_filter: self.dynamic_filter.clone(), + } + } + + /// Clone this exec, overriding only the limit hint. + pub fn with_new_limit_options(&self, limit_options: Option<LimitOptions>) -> Self { + Self { + limit_options, + // clone the rest of the fields + required_input_ordering: self.required_input_ordering.clone(), + metrics: ExecutionPlanMetricsSet::new(), + input_order_mode: self.input_order_mode.clone(), + cache: Arc::clone(&self.cache), + mode: self.mode, + group_by: Arc::clone(&self.group_by), + aggr_expr: Arc::clone(&self.aggr_expr), + filter_expr: Arc::clone(&self.filter_expr), + input: Arc::clone(&self.input), + schema: Arc::clone(&self.schema), + input_schema: Arc::clone(&self.input_schema), + dynamic_filter: self.dynamic_filter.clone(), + } + } + + pub fn cache(&self) -> &PlanProperties { + &self.cache + } + + /// Create a new hash aggregate execution plan + pub fn try_new( + mode: AggregateMode, + group_by: impl Into<Arc<PhysicalGroupBy>>, + aggr_expr: Vec<Arc<AggregateFunctionExpr>>, + filter_expr: Vec<Option<Arc<dyn PhysicalExpr>>>, + input: Arc<dyn ExecutionPlan>, + input_schema: SchemaRef, + ) -> Result<Self> { + let group_by = group_by.into(); + let schema = create_schema(&input.schema(), &group_by, &aggr_expr, mode)?; + + let schema = Arc::new(schema); + AggregateExec::try_new_with_schema( + mode, + group_by, + aggr_expr, + filter_expr, + input, + input_schema, + schema, + ) + } + + /// Create a new hash aggregate execution plan with the given schema. + /// This constructor isn't part of the public API, it is used internally + /// by DataFusion to enforce schema consistency during when re-creating + /// `AggregateExec`s inside optimization rules. Schema field names of an + /// `AggregateExec` depends on the names of aggregate expressions. Since + /// a rule may re-write aggregate expressions (e.g. reverse them) during + /// initialization, field names may change inadvertently if one re-creates + /// the schema in such cases. + fn try_new_with_schema( + mode: AggregateMode, + group_by: impl Into<Arc<PhysicalGroupBy>>, + mut aggr_expr: Vec<Arc<AggregateFunctionExpr>>, + filter_expr: impl Into<Arc<[Option<Arc<dyn PhysicalExpr>>]>>, + input: Arc<dyn ExecutionPlan>, + input_schema: SchemaRef, + schema: SchemaRef, + ) -> Result<Self> { + let group_by = group_by.into(); + let filter_expr = filter_expr.into(); + + // Make sure arguments are consistent in size + assert_eq_or_internal_err!( + aggr_expr.len(), + filter_expr.len(), + "Inconsistent aggregate expr: {:?} and filter expr: {:?} for AggregateExec, their size should match", + aggr_expr, + filter_expr + ); + + let input_eq_properties = input.equivalence_properties(); + // Get GROUP BY expressions: + let groupby_exprs = group_by.input_exprs(); + // If existing ordering satisfies a prefix of the GROUP BY expressions, + // prefix requirements with this section. In this case, aggregation will + // work more efficiently. + // Copy the `PhysicalSortExpr`s to retain the sort options. + let (new_sort_exprs, indices) = + input_eq_properties.find_longest_permutation(&groupby_exprs)?; + + let mut new_requirements = new_sort_exprs + .into_iter() + .map(PhysicalSortRequirement::from) + .collect::<Vec<_>>(); + + let req = get_finer_aggregate_exprs_requirement( + &mut aggr_expr, + &group_by, + input_eq_properties, + &mode, + )?; + new_requirements.extend(req); + + let required_input_ordering = + LexRequirement::new(new_requirements).map(OrderingRequirements::new_soft); + + // If our aggregation has grouping sets then our base grouping exprs will + // be expanded based on the flags in `group_by.groups` where for each + // group we swap the grouping expr for `null` if the flag is `true` + // That means that each index in `indices` is valid if and only if + // it is not null in every group + let indices: Vec<usize> = indices + .into_iter() + .filter(|idx| group_by.groups.iter().all(|group| !group[*idx])) + .collect(); + + let input_order_mode = if indices.len() == groupby_exprs.len() + && !indices.is_empty() + && group_by.groups.len() == 1 + { + InputOrderMode::Sorted + } else if !indices.is_empty() { + InputOrderMode::PartiallySorted(indices) + } else { + InputOrderMode::Linear + }; + + // construct a map from the input expression to the output expression of the Aggregation group by + let group_expr_mapping = + ProjectionMapping::try_new(group_by.expr.clone(), &input.schema())?; + + let cache = Self::compute_properties( + &input, + Arc::clone(&schema), + &group_expr_mapping, + group_by.is_true_no_grouping(), + &mode, + &input_order_mode, + aggr_expr.as_ref(), + )?; + + let mut exec = AggregateExec { + mode, + group_by, + aggr_expr: aggr_expr.into(), + filter_expr, + input, + schema, + input_schema, + metrics: ExecutionPlanMetricsSet::new(), + required_input_ordering, + limit_options: None, + input_order_mode, + cache: Arc::new(cache), + dynamic_filter: None, + }; + + exec.init_dynamic_filter(); + + Ok(exec) + } + + /// Aggregation mode (full, partial) + pub fn mode(&self) -> &AggregateMode { + &self.mode + } + + /// Set the limit options for this AggExec + pub fn with_limit_options(mut self, limit_options: Option<LimitOptions>) -> Self { + self.limit_options = limit_options; + self + } + + /// Get the limit options (if set) + pub fn limit_options(&self) -> Option<LimitOptions> { + self.limit_options + } + + /// Grouping expressions + pub fn group_expr(&self) -> &PhysicalGroupBy { + &self.group_by + } + + /// Grouping expressions as they occur in the output schema + pub fn output_group_expr(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.group_by.output_exprs() + } + + /// Aggregate expressions + pub fn aggr_expr(&self) -> &[Arc<AggregateFunctionExpr>] { + &self.aggr_expr + } + + /// FILTER (WHERE clause) expression for each aggregate expression + pub fn filter_expr(&self) -> &[Option<Arc<dyn PhysicalExpr>>] { + &self.filter_expr + } + + /// Returns the dynamic filter expression for this aggregate, if set. + #[deprecated( + since = "55.0.0", + note = "Use ExecutionPlan::dynamic_expressions_produced instead" + )] + pub fn dynamic_filter_expr(&self) -> Option<&Arc<DynamicFilterPhysicalExpr>> { + self.dynamic_filter.as_ref().map(|df| &df.filter) + } + + /// Replace the dynamic filter expression. This method errors if the aggregate does not + /// support dynamic filtering or if the filter expression is incompatible with this + /// [`AggregateExec`]. + pub fn with_dynamic_filter_expr( + mut self, + filter: Arc<DynamicFilterPhysicalExpr>, + ) -> Result<Self> { + // If there is no dynamic filter state initialized via `try_new`, then + // we can safely assume that the aggregate does not support dynamic filtering. + let Some(dyn_filter) = self.dynamic_filter.as_ref() else { + return internal_err!("Aggregate does not support dynamic filtering"); + }; + + // Validate that the filter is compatible with the aggregation columns. + let cols = self.cols_for_dynamic_filter(&dyn_filter.supported_accumulators_info); + if cols.len() != filter.children().len() { + return internal_err!( + "Dynamic filter expression is incompatible with aggregate due to mismatched number of columns" + ); + } + for (col, child) in cols.iter().zip(filter.children()) { + if !col.eq(child) { + return internal_err!( + "Dynamic filter expression is incompatible with aggregate due to mismatched column references {col} != {child}" + ); + } + } + + // Overwrite our filter + self.dynamic_filter = Some(Arc::new(AggrDynFilter { + filter, + supported_accumulators_info: dyn_filter.supported_accumulators_info.clone(), + })); + Ok(self) + } + + /// Input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Get the input schema before any aggregates are applied + pub fn input_schema(&self) -> SchemaRef { + Arc::clone(&self.input_schema) + } + + /// Aggregation has multiple specialized implementations optimized for + /// different workloads. This function picks the best available path. + fn execute_typed( + &self, + partition: usize, + context: &Arc<TaskContext>, + ) -> Result<StreamType> { + if self.group_by.is_true_no_grouping() { + return Ok(StreamType::AggregateStream(AggregateStream::new( + self, context, partition, + )?)); + } + + // grouping by an expression that has a sort/limit upstream + if let Some(config) = self.limit_options + && !self.is_unordered_unfiltered_group_by_distinct() + { + return Ok(StreamType::GroupedPriorityQueue( + GroupedTopKAggregateStream::new(self, context, partition, config.limit)?, + )); + } + + // Select the stream type based on the query shape and configuration. + // For an overview, see the `Aggregate planning` section in this file's + // documentation. + // + // # Implementation Note + // + // `GroupedHashAggregateStream` is being incrementally refactored. See the + // tracking issue for details. + // + // New features and improvements should go directly into the new implementation. + // Please coordinate through the tracking issue. + // + // Issue: <https://github.com/apache/datafusion/issues/22710> + if context + .session_config() + .options() + .execution + .enable_migration_aggregate + { + if self.should_use_ordered_partial_aggregate_stream(context) { + return Ok(StreamType::OrderedPartialAggregate( + OrderedPartialAggregateStream::new(self, context, partition)?, + )); + } + + if self.should_use_partial_hash_stream(context) { + return Ok(StreamType::PartialHash(PartialHashAggregateStream::new( + self, context, partition, + )?)); + } + + if self.should_use_partial_reduce_hash_stream(context) { + return Ok(StreamType::PartialReduceHash( + PartialReduceHashAggregateStream::new(self, context, partition)?, + )); + } + + if self.should_use_ordered_final_aggregate_stream(context) { + return Ok(StreamType::OrderedFinalAggregate( + OrderedFinalAggregateStream::new(self, context, partition)?, + )); + } + + if self.should_use_final_hash_stream(context) { + return Ok(StreamType::FinalHash(FinalHashAggregateStream::new( + self, context, partition, + )?)); + } + + if self.should_use_single_hash_stream(context) { + return Ok(StreamType::SingleHash(SingleHashAggregateStream::new( + self, context, partition, + )?)); + } + } + + // Execution paths that have not been migrated use the fallback implementation + Ok(StreamType::GroupedHash(GroupedHashAggregateStream::new( + self, context, partition, + )?)) + } + + fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool { + self.mode == AggregateMode::Partial + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + && self.limit_options_supported_by_hash_stream() + } + + fn should_use_ordered_partial_aggregate_stream( + &self, + _context: &TaskContext, + ) -> bool { + self.mode == AggregateMode::Partial + && self.input_order_mode != InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + && self.limit_options_supported_by_hash_stream() + } + + fn should_use_final_hash_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && self.limit_options_supported_by_hash_stream() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && 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; + } + + self.mode == AggregateMode::PartialReduce + && self.limit_options.is_none() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + fn should_use_single_hash_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + ) && self.limit_options.is_none() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + fn should_use_ordered_final_aggregate_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && self.limit_options_supported_by_hash_stream() + && self.input_order_mode != InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + /// See comments in `PartialHashAggregateStream` limit optimization section + fn limit_options_supported_by_hash_stream(&self) -> bool { + self.limit_options.is_none() || self.is_unordered_unfiltered_group_by_distinct() + } + + /// Finds the DataType and SortDirection for this Aggregate, if there is one + pub fn get_minmax_desc(&self) -> Option<(FieldRef, bool)> { + let agg_expr = self.aggr_expr.iter().exactly_one().ok()?; + agg_expr.get_minmax_desc() + } + + /// true, if this Aggregate has a group-by with no required or explicit ordering, + /// no filtering and no aggregate expressions + /// This method qualifies the use of the LimitedDistinctAggregation rewrite rule + /// on an AggregateExec. + pub fn is_unordered_unfiltered_group_by_distinct(&self) -> bool { + if self + .limit_options() + .and_then(|config| config.descending) + .is_some() + { + return false; + } + // ensure there is a group by + if self.group_expr().is_empty() && !self.group_expr().has_grouping_set() { + return false; + } + // ensure there are no aggregate expressions + if !self.aggr_expr().is_empty() { + return false; + } + // ensure there are no filters on aggregate expressions; the above check + // may preclude this case + if self.filter_expr().iter().any(|e| e.is_some()) { + return false; + } + // ensure there are no order by expressions + if !self.aggr_expr().iter().all(|e| e.order_bys().is_empty()) { + return false; + } + // ensure there is no output ordering; can this rule be relaxed? + if self.properties().output_ordering().is_some() { + return false; + } + // ensure no ordering is required on the input + if let Some(requirement) = self.required_input_ordering().swap_remove(0) { + return matches!(requirement, OrderingRequirements::Hard(_)); + } + true + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + pub fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + group_expr_mapping: &ProjectionMapping, + is_true_no_grouping: bool, + mode: &AggregateMode, + input_order_mode: &InputOrderMode, + aggr_exprs: &[Arc<AggregateFunctionExpr>], + ) -> Result<PlanProperties> { + // Construct equivalence properties: + let mut eq_properties = input + .equivalence_properties() + .project(group_expr_mapping, schema); + + // True no-group aggregates produce only one row in each output + // partition, so aggregate outputs are constants within the partition. + // Grouping sets with empty grouping expressions are not covered here: + // their output schema can include grouping-set columns before the + // aggregate columns, so this aggregate-column mapping does not apply. + if is_true_no_grouping { + let new_constants = aggr_exprs.iter().enumerate().map(|(idx, func)| { + let column = Arc::new(Column::new(func.name(), idx)); + ConstExpr::from(column as Arc<dyn PhysicalExpr>) + }); + eq_properties.add_constants(new_constants)?; + } + + // Group by expression will be a distinct value after the aggregation. + // Add it into the constraint set. + let mut constraints = eq_properties.constraints().to_vec(); + let new_constraint = Constraint::Unique( + group_expr_mapping + .iter() + .flat_map(|(_, target_cols)| { + target_cols.iter().flat_map(|(expr, _)| { + expr.downcast_ref::<Column>().map(|c| c.index()) + }) + }) + .collect(), + ); + constraints.push(new_constraint); + eq_properties = + eq_properties.with_constraints(Constraints::new_unverified(constraints)); + + // Get output partitioning: + let input_partitioning = input.output_partitioning().clone(); + let output_partitioning = match mode.input_mode() { + AggregateInputMode::Raw => { + // First stage aggregation will not change the output partitioning, + // but needs to respect aliases (e.g. mapping in the GROUP BY + // expression). + let input_eq_properties = input.equivalence_properties(); + input_partitioning.project(group_expr_mapping, input_eq_properties) + } + AggregateInputMode::Partial => input_partitioning.clone(), + }; + + // TODO: Emission type and boundedness information can be enhanced here + let emission_type = if *input_order_mode == InputOrderMode::Linear { + EmissionType::Final + } else { + input.pipeline_behavior() + }; + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + emission_type, + input.boundedness(), + )) + } + + pub fn input_order_mode(&self) -> &InputOrderMode { + &self.input_order_mode + } + + /// Estimates output statistics for this aggregate node. + /// + /// For aggregations without group-by expressions, row count follows the + /// number of logical aggregate rows and the aggregate output mode. True + /// no-group aggregates have one logical row; empty grouping sets have one + /// logical row per grouping-set occurrence. + /// + /// For grouped aggregations with known input row count > 1, the output row + /// count is estimated as: + /// + /// ```text + /// ndv = sum over each grouping set of product(max(NDV_i + nulls_i, 1)) + /// output_rows = input_rows // baseline + /// output_rows = min(output_rows, ndv) // if NDV available + /// output_rows = min(output_rows, limit) // if TopK active + /// ``` + /// + /// **Example 1 — single group key:** + /// `GROUP BY city` where input_rows = 10,000, NDV(city) = 200 + /// → output_rows = min(10_000, 200) = 200 + /// + /// **Example 2 — two group keys with TopK:** + /// `GROUP BY city, category` where input_rows = 10,000, NDV(city) = 200, + /// NDV(category) = 5, limit = 100 + /// → ndv = 200 × 5 = 1,000 + /// → output_rows = min(10_000, 1_000) = 1,000 + /// → output_rows = min(1_000, 100) = 100 + /// + /// When `input_rows` is absent but NDV is available, falls back to: + /// + /// ```text + /// output_rows = min(ndv, limit) // if both available + /// output_rows = ndv // if only NDV available + /// output_rows = limit // if only limit available + /// ``` + /// + /// NDV estimation details (see [`Self::compute_group_ndv`]): + /// - For each grouping set, only active (non-NULL) columns contribute + /// - Per-column contribution is `max(NDV + null_adj, 1)` where `null_adj` + /// is 1 when nulls are present, 0 otherwise (a null group is a distinct + /// output row; `.max(1)` prevents a zero NDV from zeroing the product) + /// - Per-set products are summed across all grouping sets + /// - Requires NDV stats for ALL active group-by columns; if any lacks stats, + /// falls back to `input_rows` (or `Absent` if that is also unknown) + fn statistics_inner( + &self, + child_statistics: &Statistics, + partition: Option<usize>, + ) -> Result<Statistics> { + // TODO stats: group expressions: + // - once expressions will be able to compute their own stats, use it here + // - case where we group by on a column for which with have the `distinct` stat + // TODO stats: aggr expression: + // - aggregations sometimes also preserve invariants such as min, max... + + let column_statistics = { + // self.schema: [<group by exprs>, <aggregate exprs>] + let mut column_statistics = Statistics::unknown_column(&self.schema()); + + for (idx, (expr, _)) in self.group_by.expr.iter().enumerate() { + if let Some(col) = expr.downcast_ref::<Column>() { + let child_col_stats = + &child_statistics.column_statistics[col.index()]; + column_statistics[idx].max_value = child_col_stats.max_value.clone(); + column_statistics[idx].min_value = child_col_stats.min_value.clone(); + column_statistics[idx].distinct_count = + child_col_stats.distinct_count; + } + } + + column_statistics + }; + match self.exact_output_rows_without_group_exprs(partition) { + Some(output_rows) => { + let total_byte_size = + Self::calculate_scaled_byte_size(child_statistics, output_rows); + + Ok(Statistics { + num_rows: Precision::Exact(output_rows), + column_statistics, + total_byte_size, + }) + } + None => { + let num_rows = self.estimate_num_rows(child_statistics, partition); + let column_statistics = self.nullify_group_columns_for_empty_input( + column_statistics, + child_statistics, + &num_rows, + ); + + let total_byte_size = num_rows + .get_value() + .and_then(|&output_rows| { + Self::calculate_scaled_byte_size(child_statistics, output_rows) + .get_value() + .map(|&bytes| Precision::Inexact(bytes)) + }) + .unwrap_or(Precision::Absent); + + Ok(Statistics { + num_rows, + column_statistics, + total_byte_size, + }) + } + } + } + + /// Exact physical output row count for aggregates without group-by + /// expressions. + /// + /// `partition` follows [`ExecutionPlan::partition_statistics`]: `Some(_)` + /// requests one output partition, while `None` requests the entire plan. + /// Partial-state output contains the logical rows in each output partition; + /// final-value output contains the global logical rows once. + /// This mirrors execution, where partial aggregation without group-by + /// expressions emits its logical rows from every output partition, including + /// empty input partitions. + /// + /// Returns `None` when grouping expressions are present and grouped + /// cardinality estimation should be used instead. + fn exact_output_rows_without_group_exprs( + &self, + partition: Option<usize>, + ) -> Option<usize> { + let logical_rows = self.logical_rows_without_group_exprs()?; + + Some(self.scale_logical_rows(logical_rows, partition)) + } + + /// Scales a logical aggregate row count to the rows this operator emits, + /// which for partial aggregation is once per output partition. + fn scale_logical_rows(&self, logical_rows: usize, partition: Option<usize>) -> usize { + match (self.mode.output_mode(), partition) { + (AggregateOutputMode::Final, _) => logical_rows, + (AggregateOutputMode::Partial, Some(_)) => logical_rows, + (AggregateOutputMode::Partial, None) => { + logical_rows * self.cache.output_partitioning().partition_count() + } + } + } + + /// Number of rows a grouped aggregate emits for an empty input. + /// + /// Grouping expressions yield no groups, so the only rows are the + /// grand-total rows of the empty grouping sets that `GROUPING SETS(())`, + /// `ROLLUP` and `CUBE` introduce alongside the non-empty ones. + fn output_rows_for_empty_input(&self, partition: Option<usize>) -> usize { + let empty_grouping_sets = self + .group_by + .groups + .iter() + .filter(|nulls| nulls.iter().all(|is_null| *is_null)) + .count(); + + self.scale_logical_rows(empty_grouping_sets, partition) + } + + /// Reports the grouping columns of an empty input as all NULL. + /// + /// The only rows such an input produces are grand-total rows, which hold + /// NULL in every grouping column, so the values copied from the child do not + /// describe the output. Rules that answer `MIN`/`MAX` from statistics read + /// these values, so an input value here becomes a wrong query result. + /// + /// The bounds are typed nulls rather than [`Precision::Absent`], both + /// because NULL is the `MIN`/`MAX` of such a column and because the data + /// type lets downstream interval analysis keep intersecting intervals of + /// that type, as `FilterExec` does for a column with no rows. + fn nullify_group_columns_for_empty_input( + &self, + mut column_statistics: Vec<ColumnStatistics>, + child_statistics: &Statistics, + num_rows: &Precision<usize>, + ) -> Vec<ColumnStatistics> { + let empty_input = child_statistics.num_rows.get_value() == Some(&0); + let emits_rows = num_rows.get_value().is_some_and(|&rows| rows > 0); + if !empty_input || !emits_rows { + return column_statistics; + } + + let schema = self.schema(); + for (idx, column_stats) in column_statistics + .iter_mut() + .take(self.group_by.expr.len()) + .enumerate() + { + let typed_null = ScalarValue::try_from(schema.field(idx).data_type()) + .unwrap_or(ScalarValue::Null); + let mut null_bound = Precision::Exact(typed_null); + if matches!(num_rows, Precision::Inexact(_)) { + null_bound = null_bound.to_inexact(); + } + column_stats.min_value = null_bound.clone(); + column_stats.max_value = null_bound; + column_stats.distinct_count = num_rows.map(|_| 0); + column_stats.null_count = *num_rows; + } + + column_statistics + } + + /// Exact number of logical aggregate rows for aggregates without group-by + /// expressions. + /// + /// A true no-group aggregate has one logical aggregate row. Empty grouping + /// sets have one logical aggregate row per grouping-set occurrence, even + /// when there are duplicate empty grouping sets. Returns `None` when there + /// are grouping expressions. + fn logical_rows_without_group_exprs(&self) -> Option<usize> { + if self.group_by.is_true_no_grouping() { + Some(1) + } else if self.group_by.expr.is_empty() { + Some(self.group_by.groups.len()) + } else { + None + } + } + + /// Estimates the output row count for grouped aggregations, combining NDV, + /// input row count, and TopK limit into a single [`Precision<usize>`]. + fn estimate_num_rows( + &self, + child_statistics: &Statistics, + partition: Option<usize>, + ) -> Precision<usize> { + let ndv = if !self.group_by.expr.is_empty() { + self.compute_group_ndv(child_statistics) + } else { + None + }; + let limit = self.limit_options.as_ref().map(|lo| lo.limit); + + if let Some(&value) = child_statistics.num_rows.get_value() { + if value > 1 { + let mut num_rows = child_statistics.num_rows.to_inexact(); + if let Some(ndv) = ndv { + num_rows = num_rows.map(|n| n.min(ndv)); + } + if let Some(limit) = limit { + num_rows = num_rows.map(|n| n.min(limit)); + } + num_rows + } else if value == 0 { + // The limit bounds groups built from input rows, not the rows + // the empty grouping sets contribute. + child_statistics + .num_rows + .map(|_| self.output_rows_for_empty_input(partition)) + } else { + let grouping_set_num = self.group_by.groups.len(); + let mut num_rows = + child_statistics.num_rows.map(|x| x * grouping_set_num); + if let Some(limit) = limit { + num_rows = num_rows.map(|n| n.min(limit)); + } + num_rows + } + } else { + match (ndv, limit) { + (Some(n), Some(l)) => Precision::Inexact(n.min(l)), + (Some(n), None) => Precision::Inexact(n), + (None, Some(l)) => Precision::Inexact(l), + (None, None) => Precision::Absent, + } + } + } + + /// Computes the estimated number of distinct groups across all grouping sets. + /// For each grouping set, computes `product(NDV_i + null_adj_i)` for active columns, + /// then sums across all sets. Returns `None` if any active column is not a direct + /// column reference or lacks `distinct_count` stats. Non-column expressions + /// (e.g. `abs(a)`) are not yet supported because expression-level statistics + /// propagation is still in progress (see <https://github.com/apache/datafusion/pull/21122>). + /// When `null_count` is absent or unknown, null_adjustment defaults to 0. + /// + /// **Single key:** `GROUP BY a` where NDV(a) = 100, null_count(a) = 5 + /// → product = max(100 + 1, 1) = 101, total = 101 + /// + /// **Two keys:** `GROUP BY a, b` where NDV(a) = 100, NDV(b) = 50, no nulls + /// → product = 100 × 50 = 5,000, total = 5,000 + /// + /// **Grouping sets:** `GROUPING SETS ((a), (b), (a, b))` with NDV(a) = 100, NDV(b) = 50 + /// → set(a) = 100, set(b) = 50, set(a, b) = 100 × 50 = 5,000 + /// → total = 100 + 50 + 5,000 = 5,150 + fn compute_group_ndv(&self, child_statistics: &Statistics) -> Option<usize> { + let mut total: usize = 0; + for group_mask in &self.group_by.groups { + let mut set_product: usize = 1; + for (j, (expr, _)) in self.group_by.expr.iter().enumerate() { + if group_mask[j] { + continue; + } + let col = expr.downcast_ref::<Column>()?; + let col_stats = &child_statistics.column_statistics[col.index()]; + let ndv = *col_stats.distinct_count.get_value()?; + let null_adjustment = match col_stats.null_count.get_value() { + Some(&n) if n > 0 => 1usize, + _ => 0, + }; + set_product = set_product + .saturating_mul(ndv.saturating_add(null_adjustment).max(1)); + } + total = total.saturating_add(set_product); + } + Some(total) + } + + /// Check if dynamic filter is possible for the current plan node. + /// - If yes, init one inside `AggregateExec`'s `dynamic_filter` field. + /// - If not supported, `self.dynamic_filter` should be kept `None` + fn init_dynamic_filter(&mut self) { + if (!self.group_by.is_empty()) || (self.mode != AggregateMode::Partial) { + debug_assert!( + self.dynamic_filter.is_none(), + "The current operator node does not support dynamic filter" + ); + return; + } + + // Already initialized. + if self.dynamic_filter.is_some() { + return; + } + + // Collect supported accumulators + // It is assumed the order of aggregate expressions are not changed from `AggregateExec` + // to `AggregateStream` + let mut aggr_dyn_filters = Vec::new(); + // All column references in the dynamic filter, used when initializing the dynamic + // filter, and it's used to decide if this dynamic filter is able to get push + // through certain node during optimization. + let mut all_cols: Vec<Arc<dyn PhysicalExpr>> = Vec::new(); + for (i, aggr_expr) in self.aggr_expr.iter().enumerate() { + // 1. Only `min` or `max` aggregate function + let fun_name = aggr_expr.fun().name(); + // HACK: Should check the function type more precisely + // Issue: <https://github.com/apache/datafusion/issues/18643> + let aggr_type = if fun_name.eq_ignore_ascii_case("min") { + DynamicFilterAggregateType::Min + } else if fun_name.eq_ignore_ascii_case("max") { + DynamicFilterAggregateType::Max + } else { + return; + }; + + // 2. arg should be only 1 column reference + if let [arg] = aggr_expr.expressions().as_slice() + && arg.is::<Column>() + { + all_cols.push(Arc::clone(arg)); + aggr_dyn_filters.push(PerAccumulatorDynFilter { + aggr_type, + aggr_index: i, + shared_bound: Arc::new(Mutex::new(ScalarValue::Null)), + }); + } + } + + if !aggr_dyn_filters.is_empty() { + self.dynamic_filter = Some(Arc::new(AggrDynFilter { + filter: Arc::new(DynamicFilterPhysicalExpr::new(all_cols, lit(true))), + supported_accumulators_info: aggr_dyn_filters, + })) + } + } + + // Collect column references for the dynamic filter expression from the supported accumulators. + fn cols_for_dynamic_filter( + &self, + supported_accumulators_info: &[PerAccumulatorDynFilter], + ) -> Vec<Arc<dyn PhysicalExpr>> { + let all_cols: Vec<Arc<dyn PhysicalExpr>> = supported_accumulators_info + .iter() + .filter_map(|info| { + // This should always be true due to how the supported accumulators + // are constructed. See `init_dynamic_filter` for more details. + if let [arg] = &self.aggr_expr[info.aggr_index].expressions().as_slice() + && arg.is::<Column>() + { + return Some(Arc::clone(arg)); + } + None + }) + .collect(); + debug_assert!(all_cols.len() == supported_accumulators_info.len()); + all_cols + } + + /// Calculate scaled byte size based on row count ratio. + /// Returns `Precision::Absent` if input statistics are insufficient. + /// Returns `Precision::Inexact` with the scaled value otherwise. + /// + /// This is a simple heuristic that assumes uniform row sizes. + #[inline] + fn calculate_scaled_byte_size( + input_stats: &Statistics, + target_row_count: usize, + ) -> Precision<usize> { + match ( + input_stats.num_rows.get_value(), + input_stats.total_byte_size.get_value(), + ) { + (Some(&input_rows), Some(&input_bytes)) if input_rows > 0 => { + let bytes_per_row = input_bytes as f64 / input_rows as f64; + let scaled_bytes = + (bytes_per_row * target_row_count as f64).ceil() as usize; + Precision::Inexact(scaled_bytes) + } + _ => Precision::Absent, + } + } +} + +impl DisplayAs for AggregateExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let format_expr_with_alias = + |(e, alias): &(Arc<dyn PhysicalExpr>, String)| -> String { + let e = e.to_string(); + if &e != alias { + format!("{e} as {alias}") + } else { + e + } + }; + + write!(f, "AggregateExec: mode={:?}", self.mode)?; + let g: Vec<String> = if self.group_by.is_single() { + self.group_by + .expr + .iter() + .map(format_expr_with_alias) + .collect() + } else { + self.group_by + .groups + .iter() + .map(|group| { + let terms = group + .iter() + .enumerate() + .map(|(idx, is_null)| { + if *is_null { + format_expr_with_alias( + &self.group_by.null_expr[idx], + ) + } else { + format_expr_with_alias(&self.group_by.expr[idx]) + } + }) + .collect::<Vec<String>>() + .join(", "); + format!("({terms})") + }) + .collect() + }; + + write!(f, ", gby=[{}]", g.join(", "))?; + + let a: Vec<String> = self + .aggr_expr + .iter() + .map(|agg| format_aggregate_exec_expr(agg).to_string()) + .collect(); + write!(f, ", aggr=[{}]", a.join(", "))?; + if let Some(config) = self.limit_options { + write!(f, ", lim=[{}]", config.limit)?; + } + + if self.input_order_mode != InputOrderMode::Linear { + write!(f, ", ordering_mode={:?}", self.input_order_mode)?; + } + } + DisplayFormatType::TreeRender => { + let format_expr_with_alias = + |(e, alias): &(Arc<dyn PhysicalExpr>, String)| -> String { + let expr_sql = fmt_sql(e.as_ref()).to_string(); + if &expr_sql != alias { + format!("{expr_sql} as {alias}") + } else { + expr_sql + } + }; + + let g: Vec<String> = if self.group_by.is_single() { + self.group_by + .expr + .iter() + .map(format_expr_with_alias) + .collect() + } else { + self.group_by + .groups + .iter() + .map(|group| { + let terms = group + .iter() + .enumerate() + .map(|(idx, is_null)| { + if *is_null { + format_expr_with_alias( + &self.group_by.null_expr[idx], + ) + } else { + format_expr_with_alias(&self.group_by.expr[idx]) + } + }) + .collect::<Vec<String>>() + .join(", "); + format!("({terms})") + }) + .collect() + }; + let a: Vec<String> = self + .aggr_expr + .iter() + .map(|agg| format_tree_aggregate_expr(agg).to_string()) + .collect(); + writeln!(f, "mode={:?}", self.mode)?; + if !g.is_empty() { + writeln!(f, "group_by={}", g.join(", "))?; + } + if !a.is_empty() { + writeln!(f, "aggr={}", a.join(", "))?; + } + if let Some(config) = self.limit_options { + writeln!(f, "limit={}", config.limit)?; + } + } + } + Ok(()) + } +} + +fn format_aggregate_exec_expr(agg: &AggregateFunctionExpr) -> Cow<'_, str> { + match agg.human_display_alias() { + Some(_) => format_human_display(agg.human_display(), agg.human_display_alias()) + .unwrap_or_else(|| Cow::Borrowed(agg.name())), + None => Cow::Borrowed(agg.name()), + } +} + +fn format_tree_aggregate_expr(agg: &AggregateFunctionExpr) -> Cow<'_, str> { + format_human_display(agg.human_display(), agg.human_display_alias()) + .unwrap_or_else(|| Cow::Borrowed(agg.name())) +} + +fn format_human_display<'a>( + human_display: Option<&'a str>, + alias: Option<&'a str>, +) -> Option<Cow<'a, str>> { + human_display.map(|human_display| match alias { + Some(alias) => Cow::Owned(format!("{human_display} as {alias}")), + None => Cow::Borrowed(human_display), + }) +} + +impl ExecutionPlan for AggregateExec { + fn name(&self) -> &'static str { + "AggregateExec" + } + + /// Return a reference to Any that can be used for down-casting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(match &self.mode { + AggregateMode::Partial | AggregateMode::PartialReduce => { + vec![Distribution::UnspecifiedDistribution] + } + AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => { + vec![Distribution::KeyPartitioned(self.group_by.input_exprs())] + } + AggregateMode::Final | AggregateMode::Single => { + vec![Distribution::SinglePartition] + } + }) + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![self.required_input_ordering.clone()] + } + + /// The output ordering of [`AggregateExec`] is determined by its `group_by` + /// columns. Although this method is not explicitly used by any optimizer + /// rules yet, overriding the default implementation ensures that it + /// accurately reflects the actual behavior. + /// + /// If the [`InputOrderMode`] is `Linear`, the `group_by` columns don't have + /// an ordering, which means the results do not either. However, in the + /// `Ordered` and `PartiallyOrdered` cases, the `group_by` columns do have + /// an ordering, which is preserved in the output. + fn maintains_input_order(&self) -> Vec<bool> { + vec![self.input_order_mode != InputOrderMode::Linear] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut me = AggregateExec::try_new_with_schema( + self.mode, + Arc::clone(&self.group_by), + self.aggr_expr.to_vec(), + Arc::clone(&self.filter_expr), + Arc::clone(&children[0]), + Arc::clone(&self.input_schema), + Arc::clone(&self.schema), + )?; + me.limit_options = self.limit_options; + me.dynamic_filter.clone_from(&self.dynamic_filter); + Ok(Arc::new(me)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let group_by = self.group_by.input_exprs(); + let aggregates = self.aggr_expr.iter().flat_map(|aggr| { + let expressions = aggr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.order_by_exprs) + }); + let filters = self.filter_expr.iter().flatten().cloned(); + let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| { + Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter.filter) + as Arc<dyn PhysicalExpr> + }); + crate::apply_expression_roots( + group_by + .into_iter() + .chain(aggregates) + .chain(filters) + .chain(dynamic_filter), + f, + ) + } + + fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.dynamic_filter + .iter() + .map(|dynamic_filter| { + Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter.filter) + as Arc<dyn PhysicalExpr> + }) + .collect() + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + self.execute_typed(partition, &context) + .map(|stream| stream.into()) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let child_statistics = Arc::clone(&input_stats[0]); + Ok(Arc::new( + self.statistics_inner(&child_statistics, args.partition())?, + )) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + /// Push down parent filters when possible (see implementation comment for details), + /// and also pushdown self dynamic filters (see `AggrDynFilter` for details) + fn gather_filters_for_pushdown( + &self, + phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + config: &ConfigOptions, + ) -> Result<FilterDescription> { + // It's safe to push down filters through aggregates when filters only reference + // grouping columns, because such filters determine which groups to compute, not + // *how* to compute them. Each group's aggregate values (SUM, COUNT, etc.) are + // calculated from the same input rows regardless of whether we filter before or + // after grouping - filtering before just eliminates entire groups early. + // This optimization is NOT safe for filters on aggregated columns (like filtering on + // the result of SUM or COUNT), as those require computing all groups first. + + // Grouping columns are output before aggregate columns, in the same order + // as the grouping expressions. A grouping-set null mask marks grouping + // columns that are not available in that set. + let mut allowed_indices: HashSet<usize> = + (0..self.group_by.expr().len()).collect(); + for null_mask in self.group_by.groups() { + allowed_indices.retain(|idx| null_mask.get(*idx) != Some(&true)); + } + + let child = self.children()[0]; + // Global aggregates and grouping sets containing an empty grouping set + // emit a row even when their input is empty. Parent filters therefore + // cannot be pushed below them, including filters without column + // references. + let may_emit_on_empty_input = self.group_by.is_true_no_grouping() + || self + .group_by + .groups() + .iter() + .any(|null_mask| null_mask.iter().all(|is_null| *is_null)); + let mut child_desc = if may_emit_on_empty_input { + ChildFilterDescription::all_unsupported(&parent_filters) + } else { + ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + allowed_indices, + child, + )? + }; + + // Include self dynamic filter when it's possible + if phase == FilterPushdownPhase::Post + && config.optimizer.enable_aggregate_dynamic_filter_pushdown + && let Some(self_dyn_filter) = &self.dynamic_filter + { + let dyn_filter = Arc::clone(&self_dyn_filter.filter); + child_desc = child_desc.with_self_filter(dyn_filter); + } + + Ok(FilterDescription::new().with_child(child_desc)) + } + + /// If child accepts self's dynamic filter, keep `self.dynamic_filter` with Some, + /// otherwise clear it to None. + fn handle_child_pushdown_result( + &self, + phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + let mut result = FilterPushdownPropagation::if_any(child_pushdown_result.clone()); + + // If this node tried to pushdown some dynamic filter before, now we check + // if the child accept the filter + if phase == FilterPushdownPhase::Post + && let Some(dyn_filter) = &self.dynamic_filter + { + let child_accepts_dyn_filter = dyn_filter + .filter + .expression_id() + .map(|id| plan_contains_expression_id(&self.input, id)) + .transpose()? + .unwrap_or(false); + + if !child_accepts_dyn_filter { + // Child can't consume the self dynamic filter, so disable it by setting + // to `None` + let mut new_node = self.clone(); + new_node.dynamic_filter = None; + + result = result + .with_updated_node(Arc::new(new_node) as Arc<dyn ExecutionPlan>); + } + } + + Ok(result) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `AggregateExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + mode, + group_by, + aggr_expr, + filter_expr, + limit_options, + input, + // Derived at construction by `create_schema` from `input_schema`, + // `group_by`, `aggr_expr` and `mode`. + schema: _, + input_schema, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + // Derived at construction from the input ordering and `group_by`. + required_input_ordering: _, + // Derived at construction from the input ordering and `group_by`. + input_order_mode: _, + // Derived at construction by `Self::compute_properties`. + cache: _, + dynamic_filter, + } = self; + + let input = ctx.encode_child(input)?; + let group_expr = + ctx.encode_expressions(group_by.expr().iter().map(|(expr, _)| expr))?; + let group_expr_name = group_by + .expr() + .iter() + .map(|(_, name)| name.to_owned()) + .collect(); + let null_expr = + ctx.encode_expressions(group_by.null_expr().iter().map(|(expr, _)| expr))?; + let groups = group_by.groups().iter().flatten().copied().collect(); + let aggr_expr_name = aggr_expr + .iter() + .map(|expr| expr.name().to_string()) + .collect(); + let aggr_expr = aggr_expr + .iter() + .map(|expr| encode_aggregate_expr(expr, ctx)) + .collect::<Result<Vec<_>>>()?; + let filter_expr = filter_expr + .iter() + .map(|filter| { + Ok(protobuf::MaybeFilter { + expr: filter + .as_ref() + .map(|expr| ctx.encode_expr(expr)) + .transpose()?, + }) + }) + .collect::<Result<Vec<_>>>()?; + // Match by name because the protobuf and execution enums use different + // discriminants, so a numeric cast would corrupt the wire format. + let mode = match mode { + AggregateMode::Partial => protobuf::AggregateMode::Partial, + AggregateMode::Final => protobuf::AggregateMode::Final, + AggregateMode::FinalPartitioned => protobuf::AggregateMode::FinalPartitioned, + AggregateMode::Single => protobuf::AggregateMode::Single, + AggregateMode::SinglePartitioned => { + protobuf::AggregateMode::SinglePartitioned + } + AggregateMode::PartialReduce => protobuf::AggregateMode::PartialReduce, + }; + let limit = limit_options.map(|options| protobuf::AggLimit { + limit: options.limit() as u64, + descending: options.descending(), + }); + // Only the shared `filter` expr is on the wire; the accumulator bounds + // in `AggrDynFilter` are runtime state repopulated during execution. + let dynamic_filter = match dynamic_filter { + Some(dynamic_filter) => { + let expr: Arc<dyn PhysicalExpr> = + Arc::clone(&dynamic_filter.filter) as Arc<dyn PhysicalExpr>; + Some(ctx.encode_expr(&expr)?) + } + None => None, + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Aggregate(Box::new( + protobuf::AggregateExecNode { + group_expr, + group_expr_name, + aggr_expr, + filter_expr, + aggr_expr_name, + mode: mode as i32, + input: Some(Box::new(input)), + input_schema: Some(input_schema.as_ref().try_into()?), + null_expr, + groups, + limit, + has_grouping_set: group_by.has_grouping_set(), + dynamic_filter, + schema: Some(self.schema.as_ref().try_into()?), + }, + )), + ), + })) + } +} + +/// Keep this marker byte-identical to the copy used by the deprecated +/// aggregate serializer in `datafusion-proto` until that path is removed. +#[cfg(feature = "proto")] +const HUMAN_DISPLAY_ALIAS_PREFIX: &str = "\u{1f}datafusion_human_display_alias_v1:"; + +#[cfg(feature = "proto")] +fn encode_human_display_alias(human_display: &str, alias: &str) -> String { + format!( + "{HUMAN_DISPLAY_ALIAS_PREFIX}{}:{alias}{human_display}", + alias.len() + ) +} + +#[cfg(feature = "proto")] +fn split_human_display_alias<'a>( + human_display: &'a str, + name: &'a str, +) -> (&'a str, Option<&'a str>) { + if let Some(encoded) = human_display.strip_prefix(HUMAN_DISPLAY_ALIAS_PREFIX) + && let Some((alias_len, encoded)) = encoded.split_once(':') + && let Ok(alias_len) = alias_len.parse::<usize>() + && let Some(alias) = encoded.get(..alias_len) + && let Some(human_display) = encoded.get(alias_len..) + && alias == name + && !human_display.is_empty() + { + return (human_display, Some(alias)); + } + + (human_display, None) +} + +#[cfg(feature = "proto")] +fn encode_aggregate_expr( + aggr_expr: &Arc<AggregateFunctionExpr>, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result<datafusion_proto_models::protobuf::PhysicalExprNode> { + use datafusion_proto_models::protobuf; + + let expressions = aggr_expr.expressions(); + let expr = ctx.encode_expressions(expressions.iter())?; + let ordering_req = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( + aggr_expr.order_bys(), + &ctx.expr_ctx(), + )?; + let name = aggr_expr.fun().name().to_string(); + // The context already applies `(!buf.is_empty()).then_some(buf)`. + let fun_definition = ctx.encode_udaf(aggr_expr.fun())?; + let human_display = match (aggr_expr.human_display(), aggr_expr.human_display_alias()) + { + (Some(display), Some(alias)) => encode_human_display_alias(display, alias), + (Some(display), None) => display.to_string(), + (None, _) => String::new(), + }; + + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::AggregateExpr( + protobuf::PhysicalAggregateExprNode { + aggregate_function: Some( + protobuf::physical_aggregate_expr_node::AggregateFunction::UserDefinedAggrFunction(name), + ), + expr, + ordering_req, + distinct: aggr_expr.is_distinct(), + ignore_nulls: aggr_expr.ignore_nulls(), + fun_definition, + human_display, + is_reversed: aggr_expr.is_reversed(), + }, + )), + }) +} + +#[cfg(feature = "proto")] +impl AggregateExec { + /// Reconstruct an [`AggregateExec`] from its protobuf representation. + /// + /// Grouping expressions are decoded against the child schema. Aggregate + /// arguments, ordering, filters, and the dynamic filter are decoded against + /// the aggregate input schema carried in the protobuf node. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_proto_models::protobuf; + use protobuf::physical_aggregate_expr_node::AggregateFunction; + use protobuf::physical_expr_node::ExprType; + + let hash_agg = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Aggregate, + "AggregateExec", + ); + // Exhaustive destructure: a new field on `AggregateExecNode` is a + // compile error here rather than a silently ignored wire field. + let protobuf::AggregateExecNode { + group_expr, + aggr_expr, + mode, + input, + group_expr_name, + aggr_expr_name, + input_schema, + null_expr, + groups, + filter_expr, + limit, + has_grouping_set, + dynamic_filter, + schema, + } = hash_agg.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "AggregateExec", "input")?; + // Match by name because the protobuf and execution enums use different + // discriminants, so a numeric cast would corrupt the wire format. + let mode = protobuf::AggregateMode::try_from(*mode).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Received an AggregateNode message with unknown AggregateMode {mode}" + ) + })?; + let mode = match mode { + protobuf::AggregateMode::Partial => AggregateMode::Partial, + protobuf::AggregateMode::Final => AggregateMode::Final, + protobuf::AggregateMode::FinalPartitioned => AggregateMode::FinalPartitioned, + protobuf::AggregateMode::Single => AggregateMode::Single, + protobuf::AggregateMode::SinglePartitioned => { + AggregateMode::SinglePartitioned + } + protobuf::AggregateMode::PartialReduce => AggregateMode::PartialReduce, + }; + let num_expr = group_expr.len(); + // Grouping expressions refer to the child plan's output schema. + let child_schema = input.schema(); + let group_expr = group_expr + .iter() + .zip(group_expr_name.iter()) + .map(|(expr, name)| { + Ok(( + ctx.decode_expr(expr, child_schema.as_ref())?, + name.to_string(), + )) + }) + .collect::<Result<Vec<_>>>()?; + let null_expr = null_expr + .iter() + .zip(group_expr_name.iter()) + .map(|(expr, name)| { + Ok(( + ctx.decode_expr(expr, child_schema.as_ref())?, + name.to_string(), + )) + }) + .collect::<Result<Vec<_>>>()?; + let groups = if groups.is_empty() { + vec![] + } else { + groups + .chunks(num_expr) + .map(|group| group.to_vec()) + .collect() + }; + // Aggregate arguments, ordering, filters, and dynamic filters refer to + // the aggregate input schema carried in the protobuf node. + let input_schema = input_schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "input_schema in AggregateNode is missing." + ) + })?; + let input_schema: SchemaRef = SchemaRef::new(input_schema.try_into()?); + let filter_expr = filter_expr + .iter() + .map(|filter| { + filter + .expr + .as_ref() + .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) + .transpose() + }) + .collect::<Result<Vec<_>>>()?; + let aggr_expr = aggr_expr + .iter() + .zip(aggr_expr_name.iter()) + .map(|(expr, name)| { + let expr_type = expr.expr_type.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "Unexpected empty aggregate physical expression" + ) + })?; + let ExprType::AggregateExpr(aggregate) = expr_type else { + return internal_err!( + "Invalid aggregate expression for AggregateExec" + ); + }; + let args = aggregate + .expr + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) + .collect::<Result<Vec<_>>>()?; + let order_by = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( + &aggregate.ordering_req, + &ctx.expr_ctx(input_schema.as_ref()), + )?; + let Some(AggregateFunction::UserDefinedAggrFunction(udaf_name)) = + aggregate.aggregate_function.as_ref() + else { + return internal_err!( + "Invalid AggregateExpr, missing aggregate_function" + ); + }; + // The context owns the payload-to-codec and + // registry-to-codec fallback order. + let udaf = + ctx.decode_udaf(udaf_name, aggregate.fun_definition.as_deref())?; + let (human_display, human_display_alias) = + split_human_display_alias(&aggregate.human_display, name); + let builder = AggregateExprBuilder::new(udaf, args) + .schema(Arc::clone(&input_schema)) + .alias(name) + .with_ignore_nulls(aggregate.ignore_nulls) + .with_distinct(aggregate.distinct) + .order_by(order_by) + .with_reversed(aggregate.is_reversed) + .human_display(human_display); + let builder = if let Some(alias) = human_display_alias { + builder.human_display_alias(alias) + } else { + builder + }; + builder.build().map(Arc::new) + }) + .collect::<Result<Vec<_>>>()?; + let group_by = + PhysicalGroupBy::new(group_expr, null_expr, groups, *has_grouping_set); + let aggregate = if let Some(schema) = schema { + let schema = SchemaRef::new(schema.try_into()?); + AggregateExec::try_new_with_schema( + mode, + group_by, + aggr_expr, + filter_expr, + input, + Arc::clone(&input_schema), + schema, + ) + } else { + AggregateExec::try_new( + mode, + group_by, + aggr_expr, + filter_expr, + input, + Arc::clone(&input_schema), + ) + }?; + let aggregate = if let Some(limit) = limit { + let options = match limit.descending { + Some(descending) => { + LimitOptions::new_with_order(limit.limit as usize, descending) + } + None => LimitOptions::new(limit.limit as usize), + }; + aggregate.with_limit_options(Some(options)) + } else { + aggregate + }; + let aggregate = if let Some(dynamic_filter) = dynamic_filter { + let dynamic_filter = + ctx.decode_expr(dynamic_filter, input_schema.as_ref())?; + let dynamic_filter = (dynamic_filter + as Arc<dyn std::any::Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + .map_err(|_| { + datafusion_common::internal_datafusion_err!( + "AggregateExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + aggregate.with_dynamic_filter_expr(dynamic_filter)? + } else { + let mut aggregate = aggregate; + aggregate.dynamic_filter = None; + aggregate + }; + + Ok(Arc::new(aggregate)) + } +} + +/// Creates the output schema for an [`AggregateExec`] containing the group by columns followed +/// by the aggregate columns. +fn create_schema( + input_schema: &Schema, + group_by: &PhysicalGroupBy, + aggr_expr: &[Arc<AggregateFunctionExpr>], + mode: AggregateMode, +) -> Result<Schema> { + let mut fields = Vec::with_capacity(group_by.num_output_exprs() + aggr_expr.len()); + fields.extend(group_by.output_fields(input_schema)?); + + match mode.output_mode() { + AggregateOutputMode::Final => { + // in final mode, the field with the final result of the accumulator + for expr in aggr_expr { + fields.push(expr.field()) + } + } + AggregateOutputMode::Partial => { + // in partial mode, the fields of the accumulator's state + for expr in aggr_expr { + fields.extend(expr.state_fields()?.iter().cloned()); + } + } + } + + Ok(Schema::new_with_metadata( + fields, + input_schema.metadata().clone(), + )) +} + +/// Determines the lexical ordering requirement for an aggregate expression. +/// +/// # Parameters +/// +/// - `aggr_expr`: A reference to an `AggregateFunctionExpr` representing the +/// aggregate expression. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// - `include_soft_requirement`: When `false`, only hard requirements are +/// considered, as indicated by [`AggregateFunctionExpr::order_sensitivity`] +/// returning [`AggregateOrderSensitivity::HardRequirement`]. +/// Otherwise, also soft requirements ([`AggregateOrderSensitivity::SoftRequirement`]) +/// are considered. +/// +/// # Returns +/// +/// A `LexOrdering` instance indicating the lexical ordering requirement for +/// the aggregate expression. +fn get_aggregate_expr_req( + aggr_expr: &AggregateFunctionExpr, + group_by: &PhysicalGroupBy, + agg_mode: &AggregateMode, + include_soft_requirement: bool, +) -> Option<LexOrdering> { + // If the aggregation is performing a "second stage" calculation, + // then ignore the ordering requirement. Ordering requirement applies + // only to the aggregation input data. + if agg_mode.input_mode() == AggregateInputMode::Partial { + return None; + } + + match aggr_expr.order_sensitivity() { + AggregateOrderSensitivity::Insensitive => return None, + AggregateOrderSensitivity::HardRequirement => {} + AggregateOrderSensitivity::SoftRequirement => { + if !include_soft_requirement { + return None; + } + } + AggregateOrderSensitivity::Beneficial => return None, + } + + let mut sort_exprs = aggr_expr.order_bys().to_vec(); + // In non-first stage modes, we accumulate data (using `merge_batch`) from + // different partitions (i.e. merge partial results). During this merge, we + // consider the ordering of each partial result. Hence, we do not need to + // use the ordering requirement in such modes as long as partial results are + // generated with the correct ordering. + if group_by.is_single() { + // Remove all orderings that occur in the group by. These requirements + // will definitely be satisfied -- Each group by expression will have + // distinct values per group, hence all requirements are satisfied. + let physical_exprs = group_by.input_exprs(); + sort_exprs.retain(|sort_expr| { + !physical_exprs_contains(&physical_exprs, &sort_expr.expr) + }); + } + LexOrdering::new(sort_exprs) +} + +/// Concatenates the given slices. +pub fn concat_slices<T: Clone>(lhs: &[T], rhs: &[T]) -> Vec<T> { + [lhs, rhs].concat() +} + +// Determines if the candidate ordering is finer than the current ordering. +// Returns `None` if they are incomparable, `Some(true)` if there is no current +// ordering or candidate ordering is finer, and `Some(false)` otherwise. +fn determine_finer( + current: &Option<LexOrdering>, + candidate: &LexOrdering, +) -> Option<bool> { + if let Some(ordering) = current { + candidate.partial_cmp(ordering).map(|cmp| cmp.is_gt()) + } else { + Some(true) + } +} + +/// Gets the common requirement that satisfies all the aggregate expressions. +/// When possible, chooses the requirement that is already satisfied by the +/// equivalence properties. +/// +/// # Parameters +/// +/// - `aggr_exprs`: A slice of `AggregateFunctionExpr` containing all the +/// aggregate expressions. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `eq_properties`: A reference to an `EquivalenceProperties` instance +/// representing equivalence properties for ordering. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// +/// # Returns +/// +/// A `Result<Vec<PhysicalSortRequirement>>` instance, which is the requirement +/// that satisfies all the aggregate requirements. Returns an error in case of +/// conflicting requirements. +pub fn get_finer_aggregate_exprs_requirement( + aggr_exprs: &mut [Arc<AggregateFunctionExpr>], + group_by: &PhysicalGroupBy, + eq_properties: &EquivalenceProperties, + agg_mode: &AggregateMode, +) -> Result<Vec<PhysicalSortRequirement>> { + let mut requirement = None; + + // First try and find a match for all hard and soft requirements. + // If a match can't be found, try a second time just matching hard + // requirements. + for include_soft_requirement in [false, true] { + for aggr_expr in aggr_exprs.iter_mut() { + let Some(aggr_req) = get_aggregate_expr_req( + aggr_expr, + group_by, + agg_mode, + include_soft_requirement, + ) + .and_then(|o| eq_properties.normalize_sort_exprs(o)) else { + // There is no aggregate ordering requirement, or it is trivially + // satisfied -- we can skip this expression. + continue; + }; + // If the common requirement is finer than the current expression's, + // we can skip this expression. If the latter is finer than the former, + // adopt it if it is satisfied by the equivalence properties. Otherwise, + // defer the analysis to the reverse expression. + let forward_finer = determine_finer(&requirement, &aggr_req); + if let Some(finer) = forward_finer { + if !finer { + continue; + } else if eq_properties.ordering_satisfy(aggr_req.clone())? { + requirement = Some(aggr_req); + continue; + } + } + if let Some(reverse_aggr_expr) = aggr_expr.reverse_expr() { + let Some(rev_aggr_req) = get_aggregate_expr_req( + &reverse_aggr_expr, + group_by, + agg_mode, + include_soft_requirement, + ) + .and_then(|o| eq_properties.normalize_sort_exprs(o)) else { + // The reverse requirement is trivially satisfied -- just reverse + // the expression and continue with the next one: + *aggr_expr = Arc::new(reverse_aggr_expr); + continue; + }; + // If the common requirement is finer than the reverse expression's, + // just reverse it and continue the loop with the next aggregate + // expression. If the latter is finer than the former, adopt it if + // it is satisfied by the equivalence properties. Otherwise, adopt + // the forward expression. + if let Some(finer) = determine_finer(&requirement, &rev_aggr_req) { + if !finer { + *aggr_expr = Arc::new(reverse_aggr_expr); + } else if eq_properties.ordering_satisfy(rev_aggr_req.clone())? { + *aggr_expr = Arc::new(reverse_aggr_expr); + requirement = Some(rev_aggr_req); + } else { + requirement = Some(aggr_req); + } + } else if forward_finer.is_some() { + requirement = Some(aggr_req); + } else { + // Neither the existing requirement nor the current aggregate + // requirement satisfy the other (forward or reverse), this + // means they are conflicting. This is a problem only for hard + // requirements. Unsatisfied soft requirements can be ignored. + if !include_soft_requirement { + return not_impl_err!( + "Conflicting ordering requirements in aggregate functions is not supported" + ); + } + } + } + } + } + + Ok(requirement.map_or_else(Vec::new, |o| o.into_iter().map(Into::into).collect())) +} + +/// Returns physical expressions for arguments to evaluate against a batch. +/// +/// The expressions are different depending on `mode`: +/// * Partial: AggregateFunctionExpr::expressions +/// * Final: columns of `AggregateFunctionExpr::state_fields()` +pub fn aggregate_expressions( + aggr_expr: &[Arc<AggregateFunctionExpr>], + mode: &AggregateMode, + col_idx_base: usize, +) -> Result<Vec<Vec<Arc<dyn PhysicalExpr>>>> { + match mode.input_mode() { + AggregateInputMode::Raw => Ok(aggr_expr + .iter() + .map(|agg| { + let mut result = agg.expressions(); + // Append ordering requirements to expressions' results. This + // way order sensitive aggregators can satisfy requirement + // themselves. + result.extend(agg.order_bys().iter().map(|item| Arc::clone(&item.expr))); + result + }) + .collect()), + AggregateInputMode::Partial => { + // In merge mode, we build the merge expressions of the aggregation. + let mut col_idx_base = col_idx_base; + aggr_expr + .iter() + .map(|agg| { + let exprs = merge_expressions(col_idx_base, agg)?; + col_idx_base += exprs.len(); + Ok(exprs) + }) + .collect() + } + } +} + +/// uses `state_fields` to build a vec of physical column expressions required to merge the +/// AggregateFunctionExpr' accumulator's state. +/// +/// `index_base` is the starting physical column index for the next expanded state field. +fn merge_expressions( + index_base: usize, + expr: &AggregateFunctionExpr, +) -> Result<Vec<Arc<dyn PhysicalExpr>>> { + expr.state_fields().map(|fields| { + fields + .iter() + .enumerate() + .map(|(idx, f)| Arc::new(Column::new(f.name(), index_base + idx)) as _) + .collect() + }) +} + +pub type AccumulatorItem = Box<dyn Accumulator>; + +pub fn create_accumulators( + aggr_expr: &[Arc<AggregateFunctionExpr>], +) -> Result<Vec<AccumulatorItem>> { + aggr_expr + .iter() + .map(|expr| expr.create_accumulator()) + .collect() +} + +/// returns a vector of ArrayRefs, where each entry corresponds to either the +/// final value (mode = Final, FinalPartitioned and Single) or states (mode = Partial) +pub fn finalize_aggregation( + accumulators: &mut [AccumulatorItem], + mode: &AggregateMode, +) -> Result<Vec<ArrayRef>> { + match mode.output_mode() { + AggregateOutputMode::Final => { + // Merge the state to the final value + accumulators + .iter_mut() + .map(|accumulator| accumulator.evaluate().and_then(|v| v.to_array())) + .collect() + } + AggregateOutputMode::Partial => { + // Build the vector of states + accumulators + .iter_mut() + .map(|accumulator| { + accumulator.state().and_then(|e| { + e.iter() + .map(|v| v.to_array()) + .collect::<Result<Vec<ArrayRef>>>() + }) + }) + .flatten_ok() + .collect() + } + } +} + +/// Evaluates groups of expressions against a record batch. +pub fn evaluate_many( + expr: &[Vec<Arc<dyn PhysicalExpr>>], + batch: &RecordBatch, +) -> Result<Vec<Vec<ArrayRef>>> { + expr.iter() + .map(|expr| evaluate_expressions_to_arrays(expr, batch)) + .collect() +} + +fn evaluate_optional( + expr: &[Option<Arc<dyn PhysicalExpr>>], + batch: &RecordBatch, +) -> Result<Vec<Option<ArrayRef>>> { + expr.iter() + .map(|expr| { + expr.as_ref() + .map(|expr| { + expr.evaluate(batch) + .and_then(|v| v.into_array(batch.num_rows())) + }) + .transpose() + }) + .collect() +} + +/// Builds the internal `__grouping_id` array for a single grouping set. +/// +/// The returned array packs two values into a single integer: +/// +/// - Low `n` bits (positions 0 .. n-1): the semantic bitmask. A `1` bit +/// at position `i` means that the `i`-th grouping column (counting from the +/// least significant bit, i.e. the *last* column in the `group` slice) is +/// `NULL` for this grouping set. +/// - High bits (positions n and above): the duplicate `ordinal`, which +/// distinguishes multiple occurrences of the same grouping-set pattern. The +/// ordinal is `0` for the first occurrence, `1` for the second, and so on. +/// +/// The integer type is chosen to be the smallest `UInt8 / UInt16 / UInt32 / +/// UInt64` that can represent both parts. It matches the type returned by +/// [`Aggregate::grouping_id_type`]. +pub(crate) fn group_id_array( + group: &[bool], + ordinal: usize, + max_ordinal: usize, + num_rows: usize, +) -> Result<ArrayRef> { + let n = group.len(); + if n > 64 { + return not_impl_err!( + "Grouping sets with more than 64 columns are not supported" + ); + } + let ordinal_bits = usize::BITS as usize - max_ordinal.leading_zeros() as usize; + let total_bits = n + ordinal_bits; + if total_bits > 64 { + return not_impl_err!( + "Grouping sets with {n} columns and a maximum duplicate ordinal of \ + {max_ordinal} require {total_bits} bits, which exceeds 64" + ); + } + let semantic_id = group.iter().fold(0u64, |acc, &is_null| { + (acc << 1) | if is_null { 1 } else { 0 } + }); + let full_id = semantic_id | ((ordinal as u64) << n); + if total_bits <= 8 { + Ok(Arc::new(UInt8Array::from(vec![full_id as u8; num_rows]))) + } else if total_bits <= 16 { + Ok(Arc::new(UInt16Array::from(vec![full_id as u16; num_rows]))) + } else if total_bits <= 32 { + Ok(Arc::new(UInt32Array::from(vec![full_id as u32; num_rows]))) + } else { + Ok(Arc::new(UInt64Array::from(vec![full_id; num_rows]))) + } +} + +/// Returns the highest duplicate ordinal across all grouping sets. +/// +/// At the call-site, the ordinal is the 0-based index assigned to each +/// occurrence of a repeated grouping-set pattern: the first occurrence gets +/// ordinal 0, the second gets 1, and so on. If the same `Vec<bool>` appears +/// three times the ordinals are 0, 1, 2 and this function returns 2. +/// Returns 0 when no grouping set is duplicated. +pub(crate) fn max_duplicate_ordinal(groups: &[Vec<bool>]) -> usize { + let mut counts: HashMap<&[bool], usize> = HashMap::new(); + for group in groups { + *counts.entry(group).or_insert(0) += 1; + } + counts.into_values().max().unwrap_or(0).saturating_sub(1) +} + +/// Evaluate a group by expression against a `RecordBatch` +/// +/// Arguments: +/// - `group_by`: the expression to evaluate +/// - `batch`: the `RecordBatch` to evaluate against +/// +/// Returns: A Vec of Vecs of Array of results +/// The outer Vec appears to be for grouping sets +/// The inner Vec contains the results per expression +/// The inner-inner Array contains the results per row +/// +/// For example, for `GROUP BY GROUPING SETS ((a, b), (a))` with input: +/// +/// ```text +/// a b +/// 1 1 +/// 1 2 +/// 2 1 +/// ``` +/// +/// The output is: +/// +/// ```text +/// [ +/// [ +/// a: [1, 1, 2] +/// b: [1, 2, 1] +/// grouping_id: [0, 0, 0] +/// ], +/// [ +/// a: [1, 1, 2] +/// b: [NULL, NULL, NULL] +/// grouping_id: [1, 1, 1] +/// ] +/// ] +/// ``` +pub fn evaluate_group_by( + group_by: &PhysicalGroupBy, + batch: &RecordBatch, +) -> Result<Vec<Vec<ArrayRef>>> { + let max_ordinal = max_duplicate_ordinal(&group_by.groups); + let mut ordinal_per_pattern: HashMap<&[bool], usize> = HashMap::new(); + let exprs = evaluate_expressions_to_arrays( + group_by.expr.iter().map(|(expr, _)| expr), + batch, + )?; + let null_exprs = evaluate_expressions_to_arrays( + group_by.null_expr.iter().map(|(expr, _)| expr), + batch, + )?; + + group_by + .groups + .iter() + .map(|group| { + let ordinal = ordinal_per_pattern.entry(group).or_insert(0); + let current_ordinal = *ordinal; + *ordinal += 1; + + let mut group_values = Vec::with_capacity(group_by.num_group_exprs()); + group_values.extend(group.iter().enumerate().map(|(idx, is_null)| { + if *is_null { + Arc::clone(&null_exprs[idx]) + } else { + Arc::clone(&exprs[idx]) + } + })); + if !group_by.is_single() { + group_values.push(group_id_array( + group, + current_ordinal, + max_ordinal, + batch.num_rows(), + )?); + } + Ok(group_values) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use std::task::{Context, Poll}; + + use super::*; + use crate::RecordBatchStream; + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::common; + use crate::common::collect; + use crate::empty::EmptyExec; + use crate::execution_plan::Boundedness; + use crate::expressions::col; + use crate::filter::FilterExecBuilder; + use crate::metrics::MetricValue; + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::test::TestMemoryExec; + use crate::test::assert_is_pending; + use crate::test::exec::{ + BlockingExec, StatisticsExec, assert_strong_count_converges_to_zero, + }; + + use arrow::array::{ + BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, + Int64Array, StructArray, UInt32Array, UInt64Array, + }; + use arrow::compute::{SortOptions, concat_batches}; + use arrow::datatypes::Int32Type; + use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; + use datafusion_common::{DataFusionError, internal_err}; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::memory_pool::FairSpillPool; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; + use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, + Signature, Volatility, + }; + use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; + use datafusion_functions_aggregate::array_agg::array_agg_udaf; + use datafusion_functions_aggregate::average::avg_udaf; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::first_last::{first_value_udaf, last_value_udaf}; + use datafusion_functions_aggregate::median::median_udaf; + use datafusion_functions_aggregate::min_max::min_udaf; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_physical_expr::Partitioning; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::{Literal, NotExpr}; + + use crate::projection::ProjectionExec; + use crate::repartition::RepartitionExec; + use datafusion_physical_expr::projection::ProjectionExpr; + use futures::{FutureExt, Stream, StreamExt}; + use insta::{allow_duplicates, assert_snapshot}; + + #[cfg(feature = "proto")] + #[test] + fn split_human_display_alias_ignores_mismatched_alias() { + let encoded = encode_human_display_alias("sum(value)", "revenue"); + + assert_eq!( + split_human_display_alias(&encoded, "other"), + (encoded.as_str(), None) + ); + } + + #[cfg(feature = "proto")] + #[test] + fn split_human_display_alias_keeps_malformed_prefix_literal() { + let display = format!("{HUMAN_DISPLAY_ALIAS_PREFIX}not-an-encoding"); + + assert_eq!( + split_human_display_alias(&display, "agg"), + (display.as_str(), None) + ); + } + + // Generate a schema which consists of 5 columns (a, b, c, d, e) + fn create_test_schema() -> Result<SchemaRef> { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, true); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, true); + let e = Field::new("e", DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); + + Ok(schema) + } + + /// some mock data to aggregates + fn some_data() -> (Arc<Schema>, Vec<RecordBatch>) { + // define a schema. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + // define data. + ( + Arc::clone(&schema), + vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![2, 3, 4, 4])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt32Array::from(vec![2, 3, 3, 4])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])), + ], + ) + .unwrap(), + ], + ) + } + + /// Generates some mock data for aggregate tests. + fn some_data_v2() -> (Arc<Schema>, Vec<RecordBatch>) { + // Define a schema: + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + // Generate data so that first and last value results are at 2nd and + // 3rd partitions. With this construction, we guarantee we don't receive + // the expected result by accident, but merging actually works properly; + // i.e. it doesn't depend on the data insertion order. + ( + Arc::clone(&schema), + vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![2, 3, 4, 4])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![2, 3, 3, 4])), + Arc::new(Float64Array::from(vec![0.0, 1.0, 2.0, 3.0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![2, 3, 3, 4])), + Arc::new(Float64Array::from(vec![3.0, 4.0, 5.0, 6.0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt32Array::from(vec![2, 3, 3, 4])), + Arc::new(Float64Array::from(vec![2.0, 3.0, 4.0, 5.0])), + ], + ) + .unwrap(), + ], + ) + } + + fn new_spill_ctx(batch_size: usize, max_memory: usize) -> Arc<TaskContext> { + let session_config = SessionConfig::new().with_batch_size(batch_size); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(max_memory))) + .build_arc() + .unwrap(); + let task_ctx = TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime); + Arc::new(task_ctx) + } + + fn migrated_hash_session_config(batch_size: usize) -> SessionConfig { + SessionConfig::new() + .with_batch_size(batch_size) + .set_bool("datafusion.execution.enable_migration_aggregate", true) + } + + fn new_migrated_hash_ctx(batch_size: usize) -> Arc<TaskContext> { + Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)), + ) + } + + fn new_finite_memory_migrated_hash_ctx( + batch_size: usize, + max_memory: usize, + ) -> Result<Arc<TaskContext>> { + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(max_memory, 1.0) + .build_arc()?; + + Ok(Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(migrated_hash_session_config(batch_size)), + )) + } + + async fn check_grouping_sets( + input: Arc<dyn ExecutionPlan>, + spill: bool, + ) -> Result<()> { + let input_schema = input.schema(); + + let grouping_set = PhysicalGroupBy::new( + vec![ + (col("a", &input_schema)?, "a".to_string()), + (col("b", &input_schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::UInt32(None)), "a".to_string()), + (lit(ScalarValue::Float64(None)), "b".to_string()), + ], + vec![ + vec![false, true], // (a, NULL) + vec![true, false], // (NULL, b) + vec![false, false], // (a,b) + ], + true, + ); + + let aggregates = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![lit(1i8)]) + .schema(Arc::clone(&input_schema)) + .alias("COUNT(1)") + .build()?, + )]; + + let task_ctx = if spill { + // adjust the max memory size to have the partial aggregate result for spill mode. + new_spill_ctx(4, 500) + } else { + Arc::new(TaskContext::default()) + }; + + let partial_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + grouping_set.clone(), + aggregates.clone(), + vec![None], + input, + Arc::clone(&input_schema), + )?); + + let result = + collect(partial_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + + if spill { + // In spill mode, we test with the limited memory, if the mem usage exceeds, + // we trigger the early emit rule, which turns out the partial aggregate result. + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&result), + @r" + +---+-----+---------------+-----------------+ + | a | b | __grouping_id | COUNT(1)[count] | + +---+-----+---------------+-----------------+ + | | 1.0 | 2 | 1 | + | | 1.0 | 2 | 1 | + | | 2.0 | 2 | 1 | + | | 2.0 | 2 | 1 | + | | 3.0 | 2 | 1 | + | | 3.0 | 2 | 1 | + | | 4.0 | 2 | 1 | + | | 4.0 | 2 | 1 | + | 2 | | 1 | 1 | + | 2 | | 1 | 1 | + | 2 | 1.0 | 0 | 1 | + | 2 | 1.0 | 0 | 1 | + | 3 | | 1 | 1 | + | 3 | | 1 | 2 | + | 3 | 2.0 | 0 | 2 | + | 3 | 3.0 | 0 | 1 | + | 4 | | 1 | 1 | + | 4 | | 1 | 2 | + | 4 | 3.0 | 0 | 1 | + | 4 | 4.0 | 0 | 2 | + +---+-----+---------------+-----------------+ + " + ); + } + } else { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&result), + @r" + +---+-----+---------------+-----------------+ + | a | b | __grouping_id | COUNT(1)[count] | + +---+-----+---------------+-----------------+ + | | 1.0 | 2 | 2 | + | | 2.0 | 2 | 2 | + | | 3.0 | 2 | 2 | + | | 4.0 | 2 | 2 | + | 2 | | 1 | 2 | + | 2 | 1.0 | 0 | 2 | + | 3 | | 1 | 3 | + | 3 | 2.0 | 0 | 2 | + | 3 | 3.0 | 0 | 1 | + | 4 | | 1 | 3 | + | 4 | 3.0 | 0 | 1 | + | 4 | 4.0 | 0 | 2 | + +---+-----+---------------+-----------------+ + " + ); + } + }; + + let merge = Arc::new(CoalescePartitionsExec::new(partial_aggregate)); + + let final_grouping_set = grouping_set.as_final(); + + let task_ctx = if spill { + new_spill_ctx(4, 3160) + } else { + task_ctx + }; + + let merged_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + final_grouping_set, + aggregates, + vec![None], + merge, + input_schema, + )?); + + let result = collect(merged_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + let batch = concat_batches(&result[0].schema(), &result)?; + assert_eq!(batch.num_columns(), 4); + assert_eq!(batch.num_rows(), 12); + + allow_duplicates! { + assert_snapshot!( + batches_to_sort_string(&result), + @r" + +---+-----+---------------+----------+ + | a | b | __grouping_id | COUNT(1) | + +---+-----+---------------+----------+ + | | 1.0 | 2 | 2 | + | | 2.0 | 2 | 2 | + | | 3.0 | 2 | 2 | + | | 4.0 | 2 | 2 | + | 2 | | 1 | 2 | + | 2 | 1.0 | 0 | 2 | + | 3 | | 1 | 3 | + | 3 | 2.0 | 0 | 2 | + | 3 | 3.0 | 0 | 1 | + | 4 | | 1 | 3 | + | 4 | 3.0 | 0 | 1 | + | 4 | 4.0 | 0 | 2 | + +---+-----+---------------+----------+ + " + ); + } + + let metrics = merged_aggregate.metrics().unwrap(); + let output_rows = metrics.output_rows().unwrap(); + assert_eq!(12, output_rows); + + Ok(()) + } + + /// build the aggregates on the data from some_data() and check the results + async fn check_aggregates(input: Arc<dyn ExecutionPlan>, spill: bool) -> Result<()> { + let input_schema = input.schema(); + + let grouping_set = PhysicalGroupBy::new( + vec![(col("a", &input_schema)?, "a".to_string())], + vec![], + vec![vec![false]], + false, + ); + + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &input_schema)?]) + .schema(Arc::clone(&input_schema)) + .alias("AVG(b)") + .build()?, + )]; + + let task_ctx = if spill { + // set to an appropriate value to trigger spill + new_spill_ctx(2, 1600) + } else { + Arc::new(TaskContext::default()) + }; + + let partial_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + grouping_set.clone(), + aggregates.clone(), + vec![None], + input, + Arc::clone(&input_schema), + )?); + + let result = + collect(partial_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + + if spill { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+---------------+-------------+ + | a | AVG(b)[count] | AVG(b)[sum] | + +---+---------------+-------------+ + | 2 | 1 | 1.0 | + | 2 | 1 | 1.0 | + | 3 | 1 | 2.0 | + | 3 | 2 | 5.0 | + | 4 | 1 | 4.0 | + | 4 | 2 | 7.0 | + +---+---------------+-------------+ + "); + } + } else { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+---------------+-------------+ + | a | AVG(b)[count] | AVG(b)[sum] | + +---+---------------+-------------+ + | 2 | 2 | 2.0 | + | 3 | 3 | 7.0 | + | 4 | 3 | 11.0 | + +---+---------------+-------------+ + "); + } + }; + + let merge = Arc::new(CoalescePartitionsExec::new(partial_aggregate)); + + let final_grouping_set = grouping_set.as_final(); + + let merged_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + final_grouping_set, + aggregates, + vec![None], + merge, + input_schema, + )?); + + // Verify statistics are preserved proportionally through aggregation + let final_stats = StatisticsContext::new() + .compute(merged_aggregate.as_ref(), &StatisticsArgs::new())?; + assert!(final_stats.total_byte_size.get_value().is_some()); + + let task_ctx = if spill { + // enlarge memory limit to let the final aggregation finish + new_spill_ctx(2, 4640) + } else { + Arc::clone(&task_ctx) + }; + let result = collect(merged_aggregate.execute(0, task_ctx)?).await?; + let batch = concat_batches(&result[0].schema(), &result)?; + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.num_rows(), 3); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+--------------------+ + | a | AVG(b) | + +---+--------------------+ + | 2 | 1.0 | + | 3 | 2.3333333333333335 | + | 4 | 3.6666666666666665 | + +---+--------------------+ + "); + // For row 2: 3, (2 + 3 + 2) / 3 + // For row 3: 4, (3 + 4 + 4) / 3 + } + + let metrics = merged_aggregate.metrics().unwrap(); + let output_rows = metrics.output_rows().unwrap(); + let spill_count = metrics.spill_count().unwrap(); + let spilled_bytes = metrics.spilled_bytes().unwrap(); + let spilled_rows = metrics.spilled_rows().unwrap(); + + assert_eq!(3, output_rows); + if spill { + assert!(spill_count > 0); + assert!(spilled_bytes > 0); + assert!(spilled_rows > 0); + } else { + assert_eq!(0, spill_count); + assert_eq!(0, spilled_bytes); + assert_eq!(0, spilled_rows); + } + + Ok(()) + } + + /// Define a test source that can yield back to runtime before returning its first item /// + + #[derive(Debug)] + struct TestYieldingExec { + /// True if this exec should yield back to runtime the first time it is polled + pub yield_first: bool, + cache: Arc<PlanProperties>, + } + + impl TestYieldingExec { + fn new(yield_first: bool) -> Self { + let schema = some_data().0; + let cache = Self::compute_properties(schema); + Self { + yield_first, + cache: Arc::new(cache), + } + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } + } + + impl DisplayAs for TestYieldingExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "TestYieldingExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } + } + + impl ExecutionPlan for TestYieldingExec { + fn name(&self) -> &'static str { + "TestYieldingExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + internal_err!("Children cannot be replaced in {self:?}") + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let stream = if self.yield_first { + TestYieldingStream::New + } else { + TestYieldingStream::Yielded + }; + + Ok(Box::pin(stream)) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if args.partition().is_some() { + return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))); + } + let (_, batches) = some_data(); + Ok(Arc::new(common::compute_record_batch_statistics( + &[batches], + &self.schema(), + None, + ))) + } + } + + /// A stream using the demo data. If inited as new, it will first yield to runtime before returning records + enum TestYieldingStream { + New, + Yielded, + ReturnedBatch1, + ReturnedBatch2, + } + + impl Stream for TestYieldingStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + match &*self { + TestYieldingStream::New => { + *(self.as_mut()) = TestYieldingStream::Yielded; + cx.waker().wake_by_ref(); + Poll::Pending + } + TestYieldingStream::Yielded => { + *(self.as_mut()) = TestYieldingStream::ReturnedBatch1; + Poll::Ready(Some(Ok(some_data().1[0].clone()))) + } + TestYieldingStream::ReturnedBatch1 => { + *(self.as_mut()) = TestYieldingStream::ReturnedBatch2; + Poll::Ready(Some(Ok(some_data().1[1].clone()))) + } + TestYieldingStream::ReturnedBatch2 => Poll::Ready(None), + } + } + } + + impl RecordBatchStream for TestYieldingStream { + fn schema(&self) -> SchemaRef { + some_data().0 + } + } + + //--- Tests ---// + + #[tokio::test] + async fn aggregate_source_not_yielding() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(false)); + + check_aggregates(input, false).await + } + + #[tokio::test] + async fn aggregate_grouping_sets_source_not_yielding() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(false)); + + check_grouping_sets(input, false).await + } + + #[tokio::test] + async fn aggregate_source_with_yielding() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(true)); + + check_aggregates(input, false).await + } + + #[tokio::test] + async fn aggregate_grouping_sets_with_yielding() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(true)); + + check_grouping_sets(input, false).await + } + + #[tokio::test] + async fn aggregate_source_not_yielding_with_spill() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(false)); + + check_aggregates(input, true).await + } + + #[tokio::test] + async fn aggregate_grouping_sets_source_not_yielding_with_spill() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(false)); + + check_grouping_sets(input, true).await + } + + #[tokio::test] + async fn aggregate_source_with_yielding_with_spill() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(true)); + + check_aggregates(input, true).await + } + + #[tokio::test] + async fn aggregate_grouping_sets_with_yielding_with_spill() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(true)); + + check_grouping_sets(input, true).await + } + + // Median(a) + fn test_median_agg_expr(schema: SchemaRef) -> Result<AggregateFunctionExpr> { + AggregateExprBuilder::new(median_udaf(), vec![col("a", &schema)?]) + .schema(schema) + .alias("MEDIAN(a)") + .build() + } + + #[tokio::test] + async fn test_oom() -> Result<()> { + let input: Arc<dyn ExecutionPlan> = Arc::new(TestYieldingExec::new(true)); + let input_schema = input.schema(); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(1, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + let groups_none = PhysicalGroupBy::default(); + let groups_some = PhysicalGroupBy::new( + vec![(col("a", &input_schema)?, "a".to_string())], + vec![], + vec![vec![false]], + false, + ); + + // something that allocates within the aggregator + let aggregates_v0: Vec<Arc<AggregateFunctionExpr>> = + vec![Arc::new(test_median_agg_expr(Arc::clone(&input_schema))?)]; + + // Use the fast path in `single_stream.rs`. + let aggregates_v2: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &input_schema)?]) + .schema(Arc::clone(&input_schema)) + .alias("AVG(b)") + .build()?, + )]; + + for (version, groups, aggregates) in [ + (0, groups_none, aggregates_v0), + (2, groups_some, aggregates_v2), + ] { + let n_aggr = aggregates.len(); + let partial_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + groups, + aggregates, + vec![None; n_aggr], + Arc::clone(&input), + Arc::clone(&input_schema), + )?); + + let stream = partial_aggregate.execute_typed(0, &task_ctx)?; + + // ensure that we really got the version we wanted + match version { + 0 => { + assert!(matches!(stream, StreamType::AggregateStream(_))); + } + 1 => { + assert!(matches!(stream, StreamType::GroupedHash(_))); + } + 2 => { + assert!(matches!(stream, StreamType::SingleHash(_))); + } + _ => panic!("Unknown version: {version}"), + } + + let stream: SendableRecordBatchStream = stream.into(); + let err = collect(stream).await.unwrap_err(); + + // error root cause traversal is a bit complicated, see #4172. + let err = err.find_root(); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "Wrong error type: {err}", + ); + } + + Ok(()) + } + + #[tokio::test] + async fn partial_grouped_aggregate_uses_raw_partial_stream() -> Result<()> { + let (schema, batches) = some_data(); + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64], + vec![DataType::Int32], + DataType::Int64, + ))); + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("input_type_asserting(b)") + .build()?, + )]; + + let partial_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + input, + Arc::clone(&schema), + )?); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(2) + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(partial_stream, StreamType::PartialHash(_))); + + let fallback_task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(2) + .set_bool("datafusion.execution.enable_migration_aggregate", false), + ), + ); + let stream = partial_aggregate.execute_typed(0, &fallback_task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + let stream: SendableRecordBatchStream = partial_stream.into(); + let batches = collect(stream).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::<Vec<_>>(), + vec![2, 1] + ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3); + + let merge = Arc::new(CoalescePartitionsExec::new(partial_aggregate)); + let final_aggregate = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggregates, + vec![None], + merge, + Arc::clone(&schema), + )?; + + let final_stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(final_stream, StreamType::FinalHash(_))); + + let stream = final_aggregate.execute_typed(0, &fallback_task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + let stream: SendableRecordBatchStream = final_stream.into(); + let batches = collect(stream).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::<Vec<_>>(), + vec![2, 1] + ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3); + + Ok(()) + } + + #[tokio::test] + async fn partial_grouped_aggregate_materializes_before_slicing() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let input_batches = vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?]; + let input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let udaf = Arc::new(AggregateUDF::from(NoFirstEmitUdaf::new())); + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("no_first_emit(value)") + .build()?, + )]; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggregates, + vec![None], + input, + Arc::clone(&schema), + )?); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(2) + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(2.0)), + ), + ), + ); + + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialHash(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let batches = collect(stream).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::<Vec<_>>(), + vec![2, 1] + ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3); + assert_snapshot!(batches_to_sort_string(&batches), @r" + +-----+-----------------------------+ + | key | no_first_emit(value)[count] | + +-----+-----------------------------+ + | 1 | 1 | + | 2 | 1 | + | 3 | 1 | + +-----+-----------------------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn limited_distinct_aggregate_uses_migrated_hash_streams() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![1, 2, 1]))], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![3, 4]))], + )?, + ]; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let partial_input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&input_batches), + Arc::clone(&schema), + None, + )?; + let partial_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + partial_input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + + let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(partial_stream, StreamType::PartialHash(_))); + let stream: SendableRecordBatchStream = partial_stream.into(); + let partial_output = collect(stream).await?; + assert_eq!( + partial_output + .iter() + .map(RecordBatch::num_rows) + .sum::<usize>(), + 2 + ); + assert_snapshot!(batches_to_sort_string(&partial_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + let final_input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + let final_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![], + vec![], + final_input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + + let final_stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(final_stream, StreamType::FinalHash(_))); + let stream: SendableRecordBatchStream = final_stream.into(); + let final_output = collect(stream).await?; + assert_eq!( + final_output + .iter() + .map(RecordBatch::num_rows) + .sum::<usize>(), + 2 + ); + assert_snapshot!(batches_to_sort_string(&final_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + Ok(()) + } + + fn single_test_aggregate() -> Result<AggregateExec> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + let input_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![input_batch]], + Arc::clone(&schema), + None, + )?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + AggregateExec::try_new( + AggregateMode::Single, + group_by, + aggregates, + vec![None], + input, + schema, + ) + } + + /// For single aggregation, ensures `SingleHashAggregateStream` is used when + /// enabled by migration config. + #[tokio::test] + async fn single_aggregate_planning() -> Result<()> { + let single = single_test_aggregate()?; + let task_ctx = new_migrated_hash_ctx(2); + + let stream = single.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 3); + assert_snapshot!(batches_to_sort_string(&output), @r" ++---+--------+ +| a | SUM(b) | ++---+--------+ +| 1 | 50.0 | +| 2 | 20.0 | +| 3 | 30.0 | ++---+--------+ +"); + + Ok(()) + } + + /// Single hash aggregation supports finite memory. + #[tokio::test] + async fn single_aggregate_with_memory_limit_planning() -> Result<()> { + let single = single_test_aggregate()?; + let task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + + let stream = single.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + + Ok(()) + } + + fn partial_reduce_test_aggregate() -> Result<AggregateExec> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates: Vec<Arc<AggregateFunctionExpr>> = 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(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let partial_reduce_input = TestMemoryExec::try_new_exec( + &[vec![partial_state_batch]], + Arc::clone(&partial_schema), + None, + )?; + + AggregateExec::try_new( + AggregateMode::PartialReduce, + group_by, + aggregates, + vec![None], + partial_reduce_input, + partial_schema, + ) + } + + /// For partial-reduce aggregation, ensures `PartialReduceHashAggregateStream` + /// is used when enabled by migration config. + #[tokio::test] + async fn partial_reduce_aggregate_planning() -> Result<()> { + let partial_reduce = partial_reduce_test_aggregate()?; + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::<usize>(), 3); + + Ok(()) + } + + /// Spilling behavior is not implemented for partial-reduce stream yet, so fall + /// back to the existing `GroupedHashAggregateStream` + #[tokio::test] + async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { + let partial_reduce = partial_reduce_test_aggregate()?; + 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), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + Ok(()) + } + + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. + #[tokio::test] + async fn ordered_partial_aggregate_planning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1])), + Arc::new(Int32Array::from(vec![10, 11, 10])), + Arc::new(Int64Array::from(vec![1, 1, 1])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2])), + Arc::new(Int32Array::from(vec![20, 21])), + Arc::new(Int64Array::from(vec![1, 1])), + ], + )?, + ]; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + + let group_by = PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value_col)") + .build()?, + )]; + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + input, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedPartialAggregate(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++----------+-----------+-------------------------+ +| sort_col | group_col | COUNT(value_col)[count] | ++----------+-----------+-------------------------+ +| 1 | 10 | 2 | +| 1 | 11 | 1 | +| 2 | 20 | 1 | +| 2 | 21 | 1 | ++----------+-----------+-------------------------+ +"); + + // Ordered partial aggregation supports finite memory. + let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + let stream = aggregate.execute_typed(0, &finite_memory_task_ctx)?; + assert!(matches!(stream, StreamType::OrderedPartialAggregate(_))); + + Ok(()) + } + + /// Ensures for ordered input, `OrderedFinalAggregateStream` is used. + #[tokio::test] + async fn ordered_final_aggregate_planning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial_aggregate = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggr_expr.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial_aggregate.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 3])), + Arc::new(Int64Array::from(vec![2, 3, 5, 7])), + ], + )?; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("key", 0), + ))]) + .unwrap(); + let final_input = + TestMemoryExec::try_new(&[vec![partial_state_batch]], partial_schema, None)? + .try_with_sort_information(vec![ordering])?; + let final_input = Arc::new(TestMemoryExec::update_cache(&Arc::new(final_input))); + + let final_aggregate = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggr_expr, + vec![None], + final_input, + Arc::clone(&schema), + )?; + assert_eq!(final_aggregate.input_order_mode(), &InputOrderMode::Sorted); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedFinalAggregate(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+--------------+ +| key | COUNT(value) | ++-----+--------------+ +| 1 | 5 | +| 2 | 5 | +| 3 | 7 | ++-----+--------------+ +"); + + // Ordered final aggregation supports finite memory. + let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + let stream = final_aggregate.execute_typed(0, &finite_memory_task_ctx)?; + assert!(matches!(stream, StreamType::OrderedFinalAggregate(_))); + + Ok(()) + } + + #[tokio::test] + async fn ordered_partial_aggregate_partially_sorted_no_emit_panic() -> Result<()> { + // Reproducer for #20445: emitting from PartiallySorted input must not + // drain more groups than the completed sort boundary allows. + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // All rows share sort_col=1, so there is no completed sort boundary + // inside this batch even though there are many distinct groups. + let n = 256; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; n])), + Arc::new(Int32Array::from((0..n as i32).collect::<Vec<_>>())), + Arc::new(Int64Array::from(vec![1; n])), + ], + )?; + + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )], + vec![None], + input, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(4096, 1.0) + .build_arc()?; + let session_config = SessionConfig::new().with_batch_size(128).set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::UInt64(Some(u64::MAX)), + ); + let task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(session_config), + ); + + let mut stream: SendableRecordBatchStream = + OrderedPartialAggregateStream::new(&aggregate, &task_ctx, 0)?.into_stream(); + + while let Some(result) = stream.next().await { + if let Err(e) = result { + if e.to_string().contains("Resources exhausted") { + break; + } + return Err(e); + } + } + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel_without_groups() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)])); + + let groups = PhysicalGroupBy::default(); + + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(a)") + .build()?, + )]; + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None], + blocking_exec, + schema, + )?); + + let fut = crate::collect(aggregate_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel_with_groups() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, true), + Field::new("b", DataType::Float64, true), + ])); + + let groups = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build()?, + )]; + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups, + aggregates.clone(), + vec![None], + blocking_exec, + schema, + )?); + + let fut = crate::collect(aggregate_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn run_first_last_multi_partitions() -> Result<()> { + for is_first_acc in [false, true] { + for spill in [false, true] { + first_last_multi_partitions(is_first_acc, spill, 5000).await? + } + } + Ok(()) + } + + // FIRST_VALUE(b ORDER BY b <SortOptions>) + fn test_first_value_agg_expr( + schema: &Schema, + sort_options: SortOptions, + ) -> Result<Arc<AggregateFunctionExpr>> { + let order_bys = vec![PhysicalSortExpr { + expr: col("b", schema)?, + options: sort_options, + }]; + let args = [col("b", schema)?]; + + AggregateExprBuilder::new(first_value_udaf(), args.to_vec()) + .order_by(order_bys) + .schema(Arc::new(schema.clone())) + .alias(String::from("first_value(b) ORDER BY [b ASC NULLS LAST]")) + .build() + .map(Arc::new) + } + + // LAST_VALUE(b ORDER BY b <SortOptions>) + fn test_last_value_agg_expr( + schema: &Schema, + sort_options: SortOptions, + ) -> Result<Arc<AggregateFunctionExpr>> { + let order_bys = vec![PhysicalSortExpr { + expr: col("b", schema)?, + options: sort_options, + }]; + let args = [col("b", schema)?]; + AggregateExprBuilder::new(last_value_udaf(), args.to_vec()) + .order_by(order_bys) + .schema(Arc::new(schema.clone())) + .alias(String::from("last_value(b) ORDER BY [b ASC NULLS LAST]")) + .build() + .map(Arc::new) + } + + fn first_value_agg_expr( + schema: &SchemaRef, + column: &str, + alias: &str, + human_display: Option<&str>, + human_display_alias: Option<&str>, + ) -> Result<AggregateFunctionExpr> { + let mut builder = + AggregateExprBuilder::new(first_value_udaf(), vec![col(column, schema)?]) + .order_by(vec![PhysicalSortExpr { + expr: col(column, schema)?, + options: SortOptions::new(false, false), + }]) + .schema(Arc::clone(schema)) + .alias(alias); + + if let Some(human_display) = human_display { + builder = builder.human_display(human_display); + } + if let Some(human_display_alias) = human_display_alias { + builder = builder.human_display_alias(human_display_alias); + } + + builder.build() + } + + #[test] + fn test_reverse_expr_preserves_aliased_human_display() -> Result<()> { + let schema = create_test_schema()?; + let agg = first_value_agg_expr( + &schema, + "b", + "agg", + Some("first_value(b) ORDER BY [b ASC NULLS LAST]"), + Some("agg"), + )?; + + let reversed = agg.reverse_expr().expect("expected reverse expr"); + + assert_eq!(reversed.name(), "agg"); + assert_eq!(reversed.human_display_alias(), Some("agg")); + assert_eq!( + format_tree_aggregate_expr(&reversed), + "last_value(b) ORDER BY [b DESC NULLS FIRST] as agg" + ); + assert_eq!( + reversed.human_display(), + Some("last_value(b) ORDER BY [b DESC NULLS FIRST]") + ); + + Ok(()) + } + + #[test] + fn test_reverse_expr_does_not_rewrite_column_names_in_human_display() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "first_value_col", + DataType::Int32, + true, + )])); + let agg = first_value_agg_expr( + &schema, + "first_value_col", + "agg", + Some( + "first_value(first_value_col) ORDER BY [first_value_col ASC NULLS LAST]", + ), + Some("agg"), + )?; + + let reversed = agg.reverse_expr().expect("expected reverse expr"); + + assert_eq!(reversed.name(), "agg"); + assert_eq!( + reversed.human_display(), + Some( + "last_value(first_value_col) ORDER BY [first_value_col DESC NULLS FIRST]" + ) + ); + assert_eq!( + format_tree_aggregate_expr(&reversed), + "last_value(first_value_col) ORDER BY [first_value_col DESC NULLS FIRST] as agg" + ); + + Ok(()) + } + + #[test] + fn test_empty_human_display_is_treated_as_absent() -> Result<()> { + let schema = create_test_schema()?; + let agg = first_value_agg_expr(&schema, "b", "agg", Some(""), None)?; + + assert_eq!(agg.human_display(), None); + assert_eq!(format_tree_aggregate_expr(&agg), "agg"); + + Ok(()) + } + + #[test] + fn test_human_display_alias_must_match_name() -> Result<()> { + let schema = create_test_schema()?; + let error = first_value_agg_expr( + &schema, + "b", + "agg", + Some("first_value(b) ORDER BY [b ASC NULLS LAST]"), + Some("other_alias"), + ) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("aggregate human_display_alias must match") + ); + + Ok(()) + } + + #[test] + fn test_reverse_expr_preserves_non_aliased_display_path() -> Result<()> { + let schema = create_test_schema()?; + let agg = first_value_agg_expr( + &schema, + "b", + "first_value(b) ORDER BY [b ASC NULLS LAST]", + None, + None, + )?; + + let reversed = agg.reverse_expr().expect("expected reverse expr"); + + assert_eq!( + reversed.name(), + "last_value(b) ORDER BY [b DESC NULLS FIRST]" + ); + assert_eq!(reversed.human_display(), None); + + Ok(()) + } + + // This function constructs the physical plan below, + // + // "AggregateExec: mode=Final, gby=[a@0 as a], aggr=[FIRST_VALUE(b)]", + // " CoalescePartitionsExec", + // " AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[FIRST_VALUE(b)], ordering_mode=None", + // " DataSourceExec: partitions=4, partition_sizes=[1, 1, 1, 1]", + // + // and checks whether the function `merge_batch` works correctly for + // FIRST_VALUE and LAST_VALUE functions. + async fn first_last_multi_partitions( + is_first_acc: bool, + spill: bool, + max_memory: usize, + ) -> Result<()> { + let task_ctx = if spill { + new_spill_ctx(2, max_memory) + } else { + Arc::new(TaskContext::default()) + }; + + let (schema, data) = some_data_v2(); + let partition1 = data[0].clone(); + let partition2 = data[1].clone(); + let partition3 = data[2].clone(); + let partition4 = data[3].clone(); + + let groups = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let sort_options = SortOptions { + descending: false, + nulls_first: false, + }; + let aggregates: Vec<Arc<AggregateFunctionExpr>> = if is_first_acc { + vec![test_first_value_agg_expr(&schema, sort_options)?] + } else { + vec![test_last_value_agg_expr(&schema, sort_options)?] + }; + + let memory_exec = TestMemoryExec::try_new_exec( + &[ + vec![partition1], + vec![partition2], + vec![partition3], + vec![partition4], + ], + Arc::clone(&schema), + None, + )?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None], + memory_exec, + Arc::clone(&schema), + )?); + let coalesce = Arc::new(CoalescePartitionsExec::new(aggregate_exec)) + as Arc<dyn ExecutionPlan>; + let aggregate_final = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + groups, + aggregates.clone(), + vec![None], + coalesce, + schema, + )?) as Arc<dyn ExecutionPlan>; + + let result = crate::collect(aggregate_final, task_ctx).await?; + if is_first_acc { + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+--------------------------------------------+ + | a | first_value(b) ORDER BY [b ASC NULLS LAST] | + +---+--------------------------------------------+ + | 2 | 0.0 | + | 3 | 1.0 | + | 4 | 3.0 | + +---+--------------------------------------------+ + "); + } + } else { + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+-------------------------------------------+ + | a | last_value(b) ORDER BY [b ASC NULLS LAST] | + +---+-------------------------------------------+ + | 2 | 3.0 | + | 3 | 5.0 | + | 4 | 6.0 | + +---+-------------------------------------------+ + "); + } + }; + Ok(()) + } + + #[tokio::test] + async fn test_get_finest_requirements() -> Result<()> { + let test_schema = create_test_schema()?; + + let options = SortOptions { + descending: false, + nulls_first: false, + }; + let col_a = &col("a", &test_schema)?; + let col_b = &col("b", &test_schema)?; + let col_c = &col("c", &test_schema)?; + let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema)); + // Columns a and b are equal. + eq_properties.add_equal_conditions(Arc::clone(col_a), Arc::clone(col_b))?; + // Aggregate requirements are + // [None], [a ASC], [a ASC, b ASC, c ASC], [a ASC, b ASC] respectively + let order_by_exprs = vec![ + vec![], + vec![PhysicalSortExpr { + expr: Arc::clone(col_a), + options, + }], + vec![ + PhysicalSortExpr { + expr: Arc::clone(col_a), + options, + }, + PhysicalSortExpr { + expr: Arc::clone(col_b), + options, + }, + PhysicalSortExpr { + expr: Arc::clone(col_c), + options, + }, + ], + vec![ + PhysicalSortExpr { + expr: Arc::clone(col_a), + options, + }, + PhysicalSortExpr { + expr: Arc::clone(col_b), + options, + }, + ], + ]; + + let common_requirement = vec![ + PhysicalSortRequirement::new(Arc::clone(col_a), Some(options)), + PhysicalSortRequirement::new(Arc::clone(col_c), Some(options)), + ]; + let mut aggr_exprs = order_by_exprs + .into_iter() + .map(|order_by_expr| { + AggregateExprBuilder::new(array_agg_udaf(), vec![Arc::clone(col_a)]) + .alias("a") + .order_by(order_by_expr) + .schema(Arc::clone(&test_schema)) + .build() + .map(Arc::new) + .unwrap() + }) + .collect::<Vec<_>>(); + let group_by = PhysicalGroupBy::new_single(vec![]); + let result = get_finer_aggregate_exprs_requirement( + &mut aggr_exprs, + &group_by, + &eq_properties, + &AggregateMode::Partial, + )?; + assert_eq!(result, common_requirement); + Ok(()) + } + + #[test] + fn test_agg_exec_same_schema() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float32, true), + Field::new("b", DataType::Float32, true), + ])); + + let col_a = col("a", &schema)?; + let option_desc = SortOptions { + descending: true, + nulls_first: true, + }; + let groups = PhysicalGroupBy::new_single(vec![(col_a, "a".to_string())]); + + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![ + test_first_value_agg_expr(&schema, option_desc)?, + test_last_value_agg_expr(&schema, option_desc)?, + ]; + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups, + aggregates, + vec![None, None], + Arc::clone(&blocking_exec) as Arc<dyn ExecutionPlan>, + schema, + )?); + let new_agg = Arc::clone(&aggregate_exec).replace_children( + vec![blocking_exec], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + assert_eq!(new_agg.schema(), aggregate_exec.schema()); + Ok(()) + } + + #[tokio::test] + async fn test_agg_exec_group_by_const() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float32, true), + Field::new("b", DataType::Float32, true), + Field::new("const", DataType::Int32, false), + ])); + + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let const_expr = Arc::new(Literal::new(ScalarValue::Int32(Some(1)))); + + let groups = PhysicalGroupBy::new( + vec![ + (col_a, "a".to_string()), + (col_b, "b".to_string()), + (const_expr, "const".to_string()), + ], + vec![ + ( + Arc::new(Literal::new(ScalarValue::Float32(None))), + "a".to_string(), + ), + ( + Arc::new(Literal::new(ScalarValue::Float32(None))), + "b".to_string(), + ), + ( + Arc::new(Literal::new(ScalarValue::Int32(None))), + "const".to_string(), + ), + ], + vec![ + vec![false, true, true], + vec![true, false, true], + vec![true, true, false], + ], + true, + ); + + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![ + AggregateExprBuilder::new(count_udaf(), vec![lit(1)]) + .schema(Arc::clone(&schema)) + .alias("1") + .build() + .map(Arc::new)?, + ]; + + let input_batches = (0..4) + .map(|_| { + let a = Arc::new(Float32Array::from(vec![0.; 8192])); + let b = Arc::new(Float32Array::from(vec![0.; 8192])); + let c = Arc::new(Int32Array::from(vec![1; 8192])); + + RecordBatch::try_new(Arc::clone(&schema), vec![a, b, c]).unwrap() + }) + .collect(); + + let input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + groups, + aggregates.clone(), + vec![None], + input, + schema, + )?); + + let output = + collect(aggregate_exec.execute(0, Arc::new(TaskContext::default()))?).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +-----+-----+-------+---------------+-------+ + | a | b | const | __grouping_id | 1 | + +-----+-----+-------+---------------+-------+ + | | | 1 | 6 | 32768 | + | | 0.0 | | 5 | 32768 | + | 0.0 | | | 3 | 32768 | + +-----+-----+-------+---------------+-------+ + "); + } + + Ok(()) + } + + #[tokio::test] + async fn test_agg_exec_struct_of_dicts() -> Result<()> { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new( + "labels".to_string(), + DataType::Struct( + vec![ + Field::new( + "a".to_string(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + ), + true, + ), + Field::new( + "b".to_string(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + ), + true, + ), + ] + .into(), + ), + false, + ), + Field::new("value", DataType::UInt64, false), + ])), + vec![ + Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new( + "a".to_string(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + ), + true, + )), + Arc::new( + vec![Some("a"), None, Some("a")] + .into_iter() + .collect::<DictionaryArray<Int32Type>>(), + ) as ArrayRef, + ), + ( + Arc::new(Field::new( + "b".to_string(), + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8), + ), + true, + )), + Arc::new( + vec![Some("b"), Some("c"), Some("b")] + .into_iter() + .collect::<DictionaryArray<Int32Type>>(), + ) as ArrayRef, + ), + ])), + Arc::new(UInt64Array::from(vec![1, 1, 1])), + ], + ) + .expect("Failed to create RecordBatch"); + + let group_by = PhysicalGroupBy::new_single(vec![( + col("labels", &batch.schema())?, + "labels".to_string(), + )]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(sum_udaf(), vec![col("value", &batch.schema())?]) + .schema(Arc::clone(&batch.schema())) + .alias(String::from("SUM(value)")) + .build() + .map(Arc::new)?, + ]; + + let input = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::<Schema>::clone(&batch.schema()), + None, + )?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + batch.schema(), + )?); + + let session_config = SessionConfig::default(); + let ctx = TaskContext::default().with_session_config(session_config); + let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +--------------+------------+ + | labels | SUM(value) | + +--------------+------------+ + | {a: a, b: b} | 2 | + | {a: , b: c} | 1 | + +--------------+------------+ + "); + } + + Ok(()) + } + + // Migrated to PartialHashAggregateStream coverage below; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_skip_aggregation_after_first_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + let input_data = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + schema, + )?); + + let mut session_config = SessionConfig::default(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(2)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); + + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let stream: SendableRecordBatchStream = Box::pin( + GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?, + ); + let output = collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 1 | + | 3 | 1 | + | 2 | 1 | + | 3 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); + } + + Ok(()) + } + + // Migrated to PartialHashAggregateStream coverage below; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_skip_aggregation_after_threshold() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + let input_data = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + schema, + )?); + + let mut session_config = SessionConfig::default(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(5)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); + + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let stream: SendableRecordBatchStream = Box::pin( + GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?, + ); + let output = collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 2 | + | 3 | 2 | + | 4 | 1 | + | 2 | 1 | + | 3 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); + } + + Ok(()) + } + + #[tokio::test] + async fn test_partial_hash_stream_skip_aggregation_after_first_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + let input_data = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + schema, + )?); + + let session_config = SessionConfig::default() + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(2)), + ) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); + + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 1 | + | 2 | 1 | + | 3 | 1 | + | 3 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); + } + + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert_eq!(skipped_rows, 3); + + Ok(()) + } + + #[tokio::test] + async fn test_partial_hash_stream_skip_aggregation_after_threshold() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + let input_data = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + schema, + )?); + + let session_config = SessionConfig::default() + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(5)), + ) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); + + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 1 | + | 2 | 2 | + | 3 | 1 | + | 3 | 2 | + | 4 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); + } + + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert_eq!(skipped_rows, 3); + + Ok(()) + } + + /// When `skip_partial_aggregation_probe_ratio_threshold` is set to 1.0, + /// the feature must be effectively disabled: even with 100% cardinality + /// (every row is a unique group), no rows should be skipped. + #[tokio::test] + async fn test_skip_aggregation_disabled_at_threshold_one() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + // Two batches are required: batch 1 triggers the probe threshold so the + // skip decision is evaluated; batch 2 is what would be skipped on main + // (where >= caused threshold=1.0 to still skip at 100% cardinality). + // All rows have unique keys => ratio = 1.0 (100% cardinality). + let input_data = vec![ + // Batch 1: fires the probe check (ratio = 5/5 = 1.0) + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0])), + ], + ) + .unwrap(), + // Batch 2: would be skipped if threshold=1.0 did not disable the feature + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![6, 7, 8, 9, 10])), + Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + schema, + )?); + + let session_config = SessionConfig::default() + .set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(1)), + ) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(1.0)), + ); + + let ctx = TaskContext::default().with_session_config(session_config); + collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + assert_eq!( + skipped_rows, 0, + "threshold=1.0 should disable skip aggregation, but {skipped_rows} rows were skipped" + ); + + Ok(()) + } + + #[test] + fn group_exprs_nullable() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float32, false), + Field::new("b", DataType::Float32, false), + ])); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("a", &input_schema)?]) + .schema(Arc::clone(&input_schema)) + .alias("COUNT(a)") + .build() + .map(Arc::new)?, + ]; + + let grouping_set = PhysicalGroupBy::new( + vec![ + (col("a", &input_schema)?, "a".to_string()), + (col("b", &input_schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Float32(None)), "a".to_string()), + (lit(ScalarValue::Float32(None)), "b".to_string()), + ], + vec![ + vec![false, true], // (a, NULL) + vec![false, false], // (a,b) + ], + true, + ); + let aggr_schema = create_schema( + &input_schema, + &grouping_set, + &aggr_expr, + AggregateMode::Final, + )?; + let expected_schema = Schema::new(vec![ + Field::new("a", DataType::Float32, false), + Field::new("b", DataType::Float32, true), + Field::new("__grouping_id", DataType::UInt8, false), + Field::new("COUNT(a)", DataType::Int64, false), + ]); + assert_eq!(aggr_schema, expected_schema); + Ok(()) + } + + // test for https://github.com/apache/datafusion/issues/13949 + async fn run_test_with_spill_pool_if_necessary( + pool_size: usize, + expect_spill: bool, + ) -> Result<()> { + fn create_record_batch( + schema: &Arc<Schema>, + data: (Vec<u32>, Vec<f64>), + ) -> Result<RecordBatch> { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(UInt32Array::from(data.0)), + Arc::new(Float64Array::from(data.1)), + ], + )?) + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + let group_keys = [2, 3, 4, 4].repeat(1_000); + let values = [1.0, 2.0, 3.0, 4.0].repeat(1_000); + let batches = vec![ + create_record_batch(&schema, (group_keys.clone(), values.clone()))?, + create_record_batch(&schema, (group_keys, values))?, + ]; + let plan: Arc<dyn ExecutionPlan> = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let grouping_set = PhysicalGroupBy::new( + vec![(col("a", &schema)?, "a".to_string())], + vec![], + vec![vec![false]], + false, + ); + + // Test with MIN for simple intermediate state (min) and AVG for multiple intermediate states (partial sum, partial count). + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![ + Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("MIN(b)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build()?, + ), + ]; + + let single_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + grouping_set, + aggregates, + vec![None, None], + plan, + Arc::clone(&schema), + )?); + + let batch_size = 2; + let memory_pool = Arc::new(FairSpillPool::new(pool_size)); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(batch_size)) + .with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(memory_pool) + .build()?, + )), + ); + + let result = collect(single_aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + + assert_spill_count_metric(expect_spill, single_aggregate); + + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+--------+--------+ + | a | MIN(b) | AVG(b) | + +---+--------+--------+ + | 2 | 1.0 | 1.0 | + | 3 | 2.0 | 2.0 | + | 4 | 3.0 | 3.5 | + +---+--------+--------+ + "); + } + + Ok(()) + } + + fn assert_spill_count_metric( + expect_spill: bool, + single_aggregate: Arc<AggregateExec>, + ) { + if let Some(metrics_set) = single_aggregate.metrics() { + let mut spill_count = 0; + + // Inspect metrics for SpillCount + for metric in metrics_set.iter() { + if let MetricValue::SpillCount(count) = metric.value() { + spill_count = count.value(); + break; + } + } + + if expect_spill && spill_count == 0 { + panic!( + "Expected spill but SpillCount metric not found or SpillCount was 0." + ); + } else if !expect_spill && spill_count > 0 { + panic!( + "Expected no spill but found SpillCount metric with value greater than 0." + ); + } + } else { + panic!("No metrics returned from the operator; cannot verify spilling."); + } + } + + #[tokio::test] + async fn test_aggregate_with_spill_if_necessary() -> Result<()> { + // test with spill + run_test_with_spill_pool_if_necessary(20_000, true).await?; + // test without spill + run_test_with_spill_pool_if_necessary(200_000, false).await?; + Ok(()) + } + + #[tokio::test] + async fn test_grouped_aggregation_respects_memory_limit() -> Result<()> { + // test with spill + fn create_record_batch( + schema: &Arc<Schema>, + data: (Vec<u32>, Vec<f64>), + ) -> Result<RecordBatch> { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(UInt32Array::from(data.0)), + Arc::new(Float64Array::from(data.1)), + ], + )?) + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + let batches = vec![ + create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?, + create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?, + ]; + let plan: Arc<dyn ExecutionPlan> = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let proj = ProjectionExec::try_new( + vec![ + ProjectionExpr::new(lit("0"), "l".to_string()), + ProjectionExpr::new_from_expression(col("a", &schema)?, &schema)?, + ProjectionExpr::new_from_expression(col("b", &schema)?, &schema)?, + ], + plan, + )?; + let plan: Arc<dyn ExecutionPlan> = Arc::new(proj); + let schema = plan.schema(); + + let grouping_set = PhysicalGroupBy::new( + vec![ + (col("l", &schema)?, "l".to_string()), + (col("a", &schema)?, "a".to_string()), + ], + vec![], + vec![vec![false, false]], + false, + ); + + // Test with MIN for simple intermediate state (min) and AVG for multiple intermediate states (partial sum, partial count). + let aggregates: Vec<Arc<AggregateFunctionExpr>> = vec![ + Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("MIN(b)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build()?, + ), + ]; + + let single_aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + grouping_set, + aggregates, + vec![None, None], + plan, + Arc::clone(&schema), + )?); + + let batch_size = 2; + let memory_pool = Arc::new(FairSpillPool::new(2000)); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(batch_size)) + .with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(memory_pool) + .build()?, + )), + ); + + let result = collect(single_aggregate.execute(0, Arc::clone(&task_ctx))?).await; + match result { + Ok(result) => { + assert_spill_count_metric(true, single_aggregate); + + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+---+--------+--------+ + | l | a | MIN(b) | AVG(b) | + +---+---+--------+--------+ + | 0 | 2 | 1.0 | 1.0 | + | 0 | 3 | 2.0 | 2.0 | + | 0 | 4 | 3.0 | 3.5 | + +---+---+--------+--------+ + "); + } + } + Err(e) => assert!(matches!(e, DataFusionError::ResourcesExhausted(_))), + } + + Ok(()) + } + + #[tokio::test] + async fn test_aggregate_statistics_edge_cases() -> Result<()> { + use datafusion_common::ColumnStatistics; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Float64, false), + ])); + + let absent_byte_stats = Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Absent, + column_statistics: vec![ + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ], + }; + let agg = build_test_aggregate( + &schema, + absent_byte_stats, + PhysicalGroupBy::default(), + None, + )?; + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + assert_eq!(stats.total_byte_size, Precision::Absent); + + let zero_row_stats = Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ], + }; + let agg_zero = build_test_aggregate( + &schema, + zero_row_stats, + PhysicalGroupBy::default(), + None, + )?; + let stats_zero = + StatisticsContext::new().compute(&agg_zero, &StatisticsArgs::new())?; + assert_eq!(stats_zero.total_byte_size, Precision::Absent); + + let single_input = + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn ExecutionPlan>; + let single_agg_zero = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::default(), + vec![count_a_aggregate(&schema)?], + vec![None], + single_input, + Arc::clone(&schema), + )?; + assert_eq!( + single_agg_zero + .properties() + .output_partitioning() + .partition_count(), + 1 + ); + let single_stats_zero = + StatisticsContext::new().compute(&single_agg_zero, &StatisticsArgs::new())?; + assert_eq!(single_stats_zero.num_rows, Precision::Exact(1)); + + Ok(()) + } + + #[tokio::test] + async fn test_aggregate_statistics_empty_input_with_grouping_sets() -> Result<()> { + let schema = empty_grouping_sets_test_schema(); + + // `GROUP BY a` produces no groups for an empty input. + let grouped = build_test_aggregate( + &schema, + empty_input_statistics(), + simple_group_by(&schema, &["a"]), + None, + )?; + let stats = StatisticsContext::new().compute(&grouped, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(0)); + + // `GROUPING SETS((a), ())`, as ROLLUP and CUBE produce, still emits the + // grand-total row of the empty grouping set on an empty input. + let with_empty_set = build_test_aggregate( + &schema, + empty_input_statistics(), + grouping_sets_with_empty(&schema, 1)?, + None, + )?; + let stats = + StatisticsContext::new().compute(&with_empty_set, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(1)); + + // `GROUPING SETS((a), (), ())` emits one grand-total row per empty + // grouping set, because execution gives each duplicate its own ordinal. + let with_duplicate_empty_sets = build_test_aggregate( + &schema, + empty_input_statistics(), + grouping_sets_with_empty(&schema, 2)?, + None, + )?; + let stats = StatisticsContext::new() + .compute(&with_duplicate_empty_sets, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(2)); + + Ok(()) + } + + /// Partial aggregation emits the grand-total row from every output + /// partition, so the whole-plan estimate scales with the partition count + /// while a single-partition request does not. + #[tokio::test] + async fn test_aggregate_statistics_empty_input_partial_mode_scaling() -> Result<()> { + let schema = empty_grouping_sets_test_schema(); + let input = Arc::new(RepartitionExec::try_new( + Arc::new(StatisticsExec::new( + empty_input_statistics(), + (*schema).clone(), + )), + Partitioning::RoundRobinBatch(4), + )?) as Arc<dyn ExecutionPlan>; + + let agg = AggregateExec::try_new( + AggregateMode::Partial, + grouping_sets_with_empty(&schema, 1)?, + vec![count_a_aggregate(&schema)?], + vec![None], + input, + Arc::clone(&schema), + )?; + assert_eq!(agg.properties().output_partitioning().partition_count(), 4); + + let context = StatisticsContext::new(); + assert_eq!( + context.compute(&agg, &StatisticsArgs::new())?.num_rows, + Precision::Exact(4) + ); + // Inexact because a repartition only estimates its per-partition row + // count. The grouping column statistics carry that same precision. + let partition_statistics = + context.compute(&agg, &StatisticsArgs::new().with_partition(Some(0)))?; + assert_eq!(partition_statistics.num_rows, Precision::Inexact(1)); + let group_column = &partition_statistics.column_statistics[0]; + let typed_null = Precision::Inexact(ScalarValue::Int32(None)); + assert_eq!(group_column.min_value, typed_null); + assert_eq!(group_column.max_value, typed_null); + assert_eq!(group_column.distinct_count, Precision::Inexact(0)); + assert_eq!(group_column.null_count, Precision::Inexact(1)); + + Ok(()) + } + + /// The input's min, max and distinct values must not reach the output + /// column statistics. See `nullify_group_columns_for_empty_input`. + #[tokio::test] + async fn test_aggregate_statistics_empty_input_nullifies_group_columns() -> Result<()> + { + let schema = empty_grouping_sets_test_schema(); + let mut input_statistics = empty_input_statistics(); + input_statistics.column_statistics[0] = ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Exact(ScalarValue::Int32(Some(5))), + min_value: Precision::Exact(ScalarValue::Int32(Some(5))), + sum_value: Precision::Absent, + distinct_count: Precision::Exact(1), + byte_size: Precision::Absent, + }; + + let agg = build_test_aggregate( + &schema, + input_statistics, + grouping_sets_with_empty(&schema, 1)?, + None, + )?; + + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(1)); + let group_column = &stats.column_statistics[0]; + let typed_null = Precision::Exact(ScalarValue::Int32(None)); + assert_eq!(group_column.min_value, typed_null); + assert_eq!(group_column.max_value, typed_null); + assert_eq!(group_column.distinct_count, Precision::Exact(0)); + assert_eq!(group_column.null_count, Precision::Exact(1)); + + Ok(()) + } + + fn empty_grouping_sets_test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Float64, false), + ])) + } + + fn empty_input_statistics() -> Statistics { + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ], + } + } + + /// `GROUPING SETS((a), (), ...)` with `empty_sets` empty grouping sets, as + /// `ROLLUP(a)` and `CUBE(a)` produce with one. + fn grouping_sets_with_empty( + schema: &SchemaRef, + empty_sets: usize, + ) -> Result<PhysicalGroupBy> { + let mut groups = vec![vec![false]]; + groups.resize(1 + empty_sets, vec![true]); + Ok(PhysicalGroupBy::new( + vec![(col("a", schema)?, "a".to_string())], + vec![(lit(ScalarValue::Int32(None)), "a".to_string())], + groups, + true, + )) + } + + fn build_test_aggregate( + schema: &SchemaRef, + stats: Statistics, + group_by: PhysicalGroupBy, + limit: Option<LimitOptions>, + ) -> Result<AggregateExec> { + build_test_aggregate_with_mode( + schema, + stats, + group_by, + limit, + AggregateMode::Final, + ) + } + + fn count_a_aggregate(schema: &SchemaRef) -> Result<Arc<AggregateFunctionExpr>> { + Ok(Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?]) + .schema(Arc::clone(schema)) + .alias("COUNT(a)") + .build()?, + )) + } + + fn build_test_aggregate_with_mode( + schema: &SchemaRef, + stats: Statistics, + group_by: PhysicalGroupBy, + limit: Option<LimitOptions>, + mode: AggregateMode, + ) -> Result<AggregateExec> { + let input = Arc::new(StatisticsExec::new(stats, (**schema).clone())) + as Arc<dyn ExecutionPlan>; + + let mut agg = AggregateExec::try_new( + mode, + group_by, + vec![count_a_aggregate(schema)?], + vec![None], + input, + Arc::clone(schema), + )?; + + if let Some(limit) = limit { + agg = agg.with_limit_options(Some(limit)); + } + + Ok(agg) + } + + fn simple_group_by(schema: &SchemaRef, cols: &[&str]) -> PhysicalGroupBy { + if cols.is_empty() { + PhysicalGroupBy::default() + } else { + PhysicalGroupBy::new_single( + cols.iter() + .map(|name| { + ( + col(name, schema).unwrap() as Arc<dyn PhysicalExpr>, + name.to_string(), + ) + }) + .collect(), + ) + } + } + + #[test] + fn test_aggregate_cardinality_estimation() -> Result<()> { + use datafusion_common::ColumnStatistics; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + + struct TestCase { + name: &'static str, + input_rows: Precision<usize>, + col_a_stats: ColumnStatistics, + col_b_stats: ColumnStatistics, + group_by_cols: Vec<&'static str>, + limit_options: Option<LimitOptions>, + expected_num_rows: Precision<usize>, + } + + let cases = vec![ + // --- NDV-based estimation --- + TestCase { + name: "single group-by col with NDV tightens estimate", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(500), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Inexact(500), + }, + TestCase { + name: "multi-col group-by multiplies NDVs", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics { + distinct_count: Precision::Exact(50), + ..ColumnStatistics::new_unknown() + }, + group_by_cols: vec!["a", "b"], + limit_options: None, + expected_num_rows: Precision::Inexact(5_000), + }, + TestCase { + name: "NDV product capped by input rows", + input_rows: Precision::Exact(200), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics { + distinct_count: Precision::Exact(50), + ..ColumnStatistics::new_unknown() + }, + group_by_cols: vec!["a", "b"], + limit_options: None, + expected_num_rows: Precision::Inexact(200), + }, + TestCase { + name: "null adjustment adds +1 per column", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(99), + null_count: Precision::Exact(10), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + // 99 + 1 (null adjustment) = 100 + expected_num_rows: Precision::Inexact(100), + }, + TestCase { + name: "null adjustment on multiple columns", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(99), + null_count: Precision::Exact(5), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics { + distinct_count: Precision::Exact(49), + null_count: Precision::Exact(3), + ..ColumnStatistics::new_unknown() + }, + group_by_cols: vec!["a", "b"], + limit_options: None, + // (99+1) * (49+1) = 100 * 50 = 5000 + expected_num_rows: Precision::Inexact(5_000), + }, + TestCase { + name: "zero null_count means no adjustment", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(100), + null_count: Precision::Exact(0), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Inexact(100), + }, + // --- Bail-out: partial NDV stats (Spark-style) --- + TestCase { + name: "bail out when one group-by col lacks NDV", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a", "b"], + limit_options: None, + expected_num_rows: Precision::Inexact(1_000_000), + }, + TestCase { + name: "bail out when all group-by cols lack NDV", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics::new_unknown(), + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Inexact(1_000_000), + }, + // --- TopK limit capping --- + TestCase { + name: "TopK limit caps output rows", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics::new_unknown(), + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: Some(LimitOptions::new(10)), + expected_num_rows: Precision::Inexact(10), + }, + TestCase { + name: "NDV + TopK limit: min(NDV, limit) when NDV < limit", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(5), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: Some(LimitOptions::new(10)), + expected_num_rows: Precision::Inexact(5), + }, + TestCase { + name: "NDV + TopK limit: min(NDV, limit) when limit < NDV", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(500), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: Some(LimitOptions::new(10)), + expected_num_rows: Precision::Inexact(10), + }, + // --- Absent input rows --- + TestCase { + name: "absent input rows without limit stays absent", + input_rows: Precision::Absent, + col_a_stats: ColumnStatistics::new_unknown(), + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Absent, + }, + TestCase { + name: "absent input rows with TopK limit gives inexact(limit)", + input_rows: Precision::Absent, + col_a_stats: ColumnStatistics::new_unknown(), + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: Some(LimitOptions::new(10)), + expected_num_rows: Precision::Inexact(10), + }, + // --- No group-by (global aggregation) --- + TestCase { + name: "no group-by cols (Final mode) returns Exact(1)", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics::new_unknown(), + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec![], + limit_options: None, + expected_num_rows: Precision::Exact(1), + }, + // --- One input row --- + TestCase { + name: "one input row returns Exact(1)", + input_rows: Precision::Exact(1), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(1), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Exact(1), + }, + // --- Zero input rows --- + TestCase { + name: "zero input rows returns Exact(0)", + input_rows: Precision::Exact(0), + col_a_stats: ColumnStatistics::new_unknown(), + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Exact(0), + }, + // --- Inexact NDV stats --- + TestCase { + name: "inexact NDV still used for estimation", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Inexact(200), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Inexact(200), + }, + TestCase { + name: "inexact NDV combined with limit", + input_rows: Precision::Exact(1_000_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Inexact(200), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: Some(LimitOptions::new(10)), + expected_num_rows: Precision::Inexact(10), + }, + // --- NDV zero column (all-null) --- + TestCase { + name: "all-null column contributes 1 to the product, not 0", + input_rows: Precision::Exact(1_000), + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(1_000), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics { + distinct_count: Precision::Exact(50), + ..ColumnStatistics::new_unknown() + }, + group_by_cols: vec!["a", "b"], + limit_options: None, + // NDV(a)=0 with nulls => max(0+1, 1)=1, NDV(b)=50 => 1*50=50 + expected_num_rows: Precision::Inexact(50), + }, + // --- Absent num_rows with NDV --- + TestCase { + name: "absent num_rows falls back to NDV estimate", + input_rows: Precision::Absent, + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: None, + expected_num_rows: Precision::Inexact(100), + }, + TestCase { + name: "absent num_rows with NDV and limit returns min(ndv, limit)", + input_rows: Precision::Absent, + col_a_stats: ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + col_b_stats: ColumnStatistics::new_unknown(), + group_by_cols: vec!["a"], + limit_options: Some(LimitOptions::new(10)), + expected_num_rows: Precision::Inexact(10), + }, + ]; + + for case in cases { + let input_stats = Statistics { + num_rows: case.input_rows, + total_byte_size: Precision::Inexact(1_000_000), + column_statistics: vec![ + case.col_a_stats.clone(), + case.col_b_stats.clone(), + ], + }; + + let group_by = simple_group_by(&schema, &case.group_by_cols); + let agg = + build_test_aggregate(&schema, input_stats, group_by, case.limit_options)?; + + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + assert_eq!( + stats.num_rows, case.expected_num_rows, + "FAILED: '{}' — expected {:?}, got {:?}", + case.name, case.expected_num_rows, stats.num_rows + ); + } + + Ok(()) + } + + #[test] + fn test_aggregate_stats_distinct_count_propagation() -> Result<()> { + use datafusion_common::ColumnStatistics; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + + let input_stats = Statistics { + num_rows: Precision::Exact(1000), + total_byte_size: Precision::Inexact(10000), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(100), + null_count: Precision::Exact(5), + ..ColumnStatistics::new_unknown() + }, + ColumnStatistics::new_unknown(), + ], + }; + let agg = build_test_aggregate( + &schema, + input_stats, + simple_group_by(&schema, &["a"]), + None, + )?; + + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + assert_eq!( + stats.column_statistics[0].distinct_count, + Precision::Exact(100), + "distinct_count should be propagated from child for group-by columns" + ); + + Ok(()) + } + + #[test] + fn test_aggregate_stats_grouping_sets() -> Result<()> { + use datafusion_common::ColumnStatistics; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + + let input_stats = Statistics { + num_rows: Precision::Exact(1_000_000), + total_byte_size: Precision::Inexact(1_000_000), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + ColumnStatistics { + distinct_count: Precision::Exact(50), + ..ColumnStatistics::new_unknown() + }, + ], + }; + + // CUBE-like grouping set: (a, NULL), (NULL, b), (a, b) — 3 groups + let grouping_set = PhysicalGroupBy::new( + vec![ + (col("a", &schema)? as Arc<dyn PhysicalExpr>, "a".to_string()), + (col("b", &schema)? as Arc<dyn PhysicalExpr>, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Int32(None)), "a".to_string()), + (lit(ScalarValue::Int32(None)), "b".to_string()), + ], + vec![ + vec![false, true], // (a, NULL) + vec![true, false], // (NULL, b) + vec![false, false], // (a, b) + ], + true, + ); + + let agg = build_test_aggregate(&schema, input_stats, grouping_set, None)?; + + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + // Per-set NDV: (a,NULL)=100, (NULL,b)=50, (a,b)=100*50=5000 + // Total = 100 + 50 + 5000 = 5150 + assert_eq!( + stats.num_rows, + Precision::Inexact(5_150), + "grouping sets should sum per-set NDV products" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_aggregate_stats_duplicate_empty_grouping_sets() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let duplicate_empty_grouping_sets = + PhysicalGroupBy::new(vec![], vec![], vec![vec![], vec![]], true); + + let single_input = + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn ExecutionPlan>; + let single_agg = AggregateExec::try_new( + AggregateMode::Single, + duplicate_empty_grouping_sets.clone(), + vec![count_a_aggregate(&schema)?], + vec![None], + single_input, + Arc::clone(&schema), + )?; + assert_eq!( + StatisticsContext::new() + .compute(&single_agg, &StatisticsArgs::new())? + .num_rows, + Precision::Exact(2) + ); + + let partial_input = + Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(2)) + as Arc<dyn ExecutionPlan>; + let partial_agg = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + duplicate_empty_grouping_sets, + vec![count_a_aggregate(&schema)?], + vec![None], + partial_input, + Arc::clone(&schema), + )?); + + assert_eq!( + partial_agg + .properties() + .output_partitioning() + .partition_count(), + 2 + ); + let task_ctx = Arc::new(TaskContext::default()); + for partition in 0..2 { + assert_eq!( + StatisticsContext::new() + .compute( + partial_agg.as_ref(), + &StatisticsArgs::new().with_partition(Some(partition)), + )? + .num_rows, + Precision::Exact(2) + ); + let result = + collect(partial_agg.execute(partition, Arc::clone(&task_ctx))?).await?; + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::<usize>(), 2); + } + + assert_eq!( + StatisticsContext::new() + .compute(partial_agg.as_ref(), &StatisticsArgs::new())? + .num_rows, + Precision::Exact(4) + ); + + Ok(()) + } + + #[test] + fn test_aggregate_stats_non_column_expr_bails_out() -> Result<()> { + use datafusion_common::ColumnStatistics; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::BinaryExpr; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + + let input_stats = Statistics { + num_rows: Precision::Exact(1_000_000), + total_byte_size: Precision::Inexact(1_000_000), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(100), + ..ColumnStatistics::new_unknown() + }, + ColumnStatistics { + distinct_count: Precision::Exact(50), + ..ColumnStatistics::new_unknown() + }, + ], + }; + + // GROUP BY (a + b) — not a direct column reference + let expr_a_plus_b: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new( + col("a", &schema)?, + Operator::Plus, + col("b", &schema)?, + )); + + let group_by = + PhysicalGroupBy::new_single(vec![(expr_a_plus_b, "a+b".to_string())]); + let agg = build_test_aggregate(&schema, input_stats, group_by, None)?; + + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + assert_eq!( + stats.num_rows, + Precision::Inexact(1_000_000), + "non-column group-by expression should bail out to input_rows" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_order_is_retained_when_spilling() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + Field::new("c", DataType::Int64, false), + ])); + + let batches = vec![vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![2])), + Arc::new(Int64Array::from(vec![2])), + Arc::new(Int64Array::from(vec![1])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(Int64Array::from(vec![1])), + Arc::new(Int64Array::from(vec![1])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![0])), + Arc::new(Int64Array::from(vec![0])), + Arc::new(Int64Array::from(vec![1])), + ], + )?, + ]]; + let scan = TestMemoryExec::try_new(&batches, Arc::clone(&schema), None)?; + let scan = scan.try_with_sort_information(vec![ + LexOrdering::new([PhysicalSortExpr::new( + col("b", schema.as_ref())?, + SortOptions::default().desc(), + )]) + .unwrap(), + ])?; + + let aggr = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new( + vec![ + (col("b", schema.as_ref())?, "b".to_string()), + (col("c", schema.as_ref())?, "c".to_string()), + ], + vec![], + vec![vec![false, false]], + false, + ), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("c", schema.as_ref())?]) + .schema(Arc::clone(&schema)) + .alias("SUM(c)") + .build()?, + )], + vec![None], + Arc::new(scan) as Arc<dyn ExecutionPlan>, + Arc::clone(&schema), + )?); + + let task_ctx = new_spill_ctx(1, 600); + let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await?; + assert_spill_count_metric(true, aggr); + + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+---+--------+ + | b | c | SUM(c) | + +---+---+--------+ + | 2 | 1 | 1 | + | 1 | 1 | 1 | + | 0 | 1 | 1 | + +---+---+--------+ + "); + } + Ok(()) + } + + /// Tests that when the memory pool is too small to accommodate the sort + /// reservation during spill, the error is properly propagated as + /// ResourcesExhausted rather than silently exceeding memory limits. + #[tokio::test] + async fn test_sort_reservation_fails_during_spill() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + Field::new("c", DataType::Float64, false), + Field::new("d", DataType::Float64, false), + Field::new("e", DataType::Float64, false), + ])); + + let batches = vec![vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1])), + Arc::new(Float64Array::from(vec![10.0])), + Arc::new(Float64Array::from(vec![20.0])), + Arc::new(Float64Array::from(vec![30.0])), + Arc::new(Float64Array::from(vec![40.0])), + Arc::new(Float64Array::from(vec![50.0])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![2])), + Arc::new(Float64Array::from(vec![11.0])), + Arc::new(Float64Array::from(vec![21.0])), + Arc::new(Float64Array::from(vec![31.0])), + Arc::new(Float64Array::from(vec![41.0])), + Arc::new(Float64Array::from(vec![51.0])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![3])), + Arc::new(Float64Array::from(vec![12.0])), + Arc::new(Float64Array::from(vec![22.0])), + Arc::new(Float64Array::from(vec![32.0])), + Arc::new(Float64Array::from(vec![42.0])), + Arc::new(Float64Array::from(vec![52.0])), + ], + )?, + ]]; + + let scan = TestMemoryExec::try_new(&batches, Arc::clone(&schema), None)?; + + let aggr = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new( + vec![(col("g", schema.as_ref())?, "g".to_string())], + vec![], + vec![vec![false]], + false, + ), + vec![ + Arc::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![col("a", schema.as_ref())?], + ) + .schema(Arc::clone(&schema)) + .alias("AVG(a)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![col("b", schema.as_ref())?], + ) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![col("c", schema.as_ref())?], + ) + .schema(Arc::clone(&schema)) + .alias("AVG(c)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![col("d", schema.as_ref())?], + ) + .schema(Arc::clone(&schema)) + .alias("AVG(d)") + .build()?, + ), + Arc::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![col("e", schema.as_ref())?], + ) + .schema(Arc::clone(&schema)) + .alias("AVG(e)") + .build()?, + ), + ], + vec![None, None, None, None, None], + Arc::new(scan) as Arc<dyn ExecutionPlan>, + Arc::clone(&schema), + )?); + + // Pool must be large enough for accumulation to start but too small for + // sort_memory after clearing. + let task_ctx = new_spill_ctx(1, 500); + let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await; + + match &result { + Ok(_) => panic!("Expected ResourcesExhausted error but query succeeded"), + Err(e) => { + let root = e.find_root(); + assert!( + matches!(root, DataFusionError::ResourcesExhausted(_)), + "Expected ResourcesExhausted, got: {root}", + ); + } + } + + Ok(()) + } + + /// Tests that PartialReduce mode: + /// 1. Accepts state as input (like Final) + /// 2. Produces state as output (like Partial) + /// 3. Can be followed by a Final stage to get the correct result + /// + /// This simulates a tree-reduce pattern: + /// Partial -> PartialReduce -> Final + async fn evaluate_partial_reduce( + groups: PhysicalGroupBy, + aggregates: Vec<Arc<AggregateFunctionExpr>>, + partition_1_and_2_batches: [Vec<RecordBatch>; 2], + ) -> Result<Vec<RecordBatch>> { + let schema = partition_1_and_2_batches + .iter() + .flatten() + .next() + .expect("Must have at least 1 batch") + .schema(); + + let [partition_1, partition_2] = partition_1_and_2_batches; + + // Step 1: Partial aggregation on partition 1 + let input1 = + TestMemoryExec::try_new_exec(&[partition_1], Arc::clone(&schema), None)?; + let partial1 = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + input1, + Arc::clone(&schema), + )?); + + // Step 2: Partial aggregation on partition 2 + let input2 = + TestMemoryExec::try_new_exec(&[partition_2], Arc::clone(&schema), None)?; + let partial2 = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + input2, + Arc::clone(&schema), + )?); + + // Collect partial results + let task_ctx = Arc::new(TaskContext::default()); + let partial_result1 = + crate::collect(Arc::clone(&partial1) as _, Arc::clone(&task_ctx)).await?; + let partial_result2 = + crate::collect(Arc::clone(&partial2) as _, Arc::clone(&task_ctx)).await?; + + // The partial results have state schema (group cols + accumulator state) + let partial_schema = partial1.schema(); + + // Step 3: PartialReduce — combine partial results, still producing state + let combined_input = TestMemoryExec::try_new_exec( + &[partial_result1, partial_result2], + Arc::clone(&partial_schema), + None, + )?; + // Coalesce into a single partition for the PartialReduce + let coalesced = Arc::new(CoalescePartitionsExec::new(combined_input)); + + let partial_reduce = Arc::new(AggregateExec::try_new( + AggregateMode::PartialReduce, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + coalesced, + Arc::clone(&partial_schema), + )?); + + // Verify PartialReduce output schema matches Partial output schema + // (both produce state, not final values) + assert_eq!(partial_reduce.schema(), partial_schema); + + // Collect PartialReduce results + let reduce_result = + crate::collect(Arc::clone(&partial_reduce) as _, Arc::clone(&task_ctx)) + .await?; + + // Step 4: Final aggregation on the PartialReduce output + let final_input = TestMemoryExec::try_new_exec( + &[reduce_result], + Arc::clone(&partial_schema), + None, + )?; + let final_agg = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + final_input, + Arc::clone(&partial_schema), + )?); + + let result = crate::collect(final_agg, Arc::clone(&task_ctx)).await?; + + Ok(result) + } + + /// Builds the shared `Partial -> PartialReduce -> Final` fixture used by + /// the `test_partial_reduce_*` tests below and runs the pipeline against + /// the aggregate produced by `build_aggregates`. + /// + /// Each test only needs to supply the UDAF/alias under test, so the test + /// body stays focused on which aggregate shape is being exercised. + async fn run_partial_reduce_pipeline<F>( + build_aggregates: F, + ) -> Result<Vec<RecordBatch>> + where + F: FnOnce(&Arc<Schema>) -> Result<Vec<Arc<AggregateFunctionExpr>>>, + { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + // Two partitions of input data so the Partial stage produces multiple + // partial states that PartialReduce must combine. + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])), + ], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3])), + Arc::new(Float64Array::from(vec![40.0, 50.0, 60.0])), + ], + )?; + + let groups = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates = build_aggregates(&schema)?; + + evaluate_partial_reduce(groups, aggregates, [vec![batch1], vec![batch2]]).await + } + + // ------------------------------------------------------------------- + // PartialReduce regression coverage. + // + // Each shape (single state field / single input arg, multi-state / + // single-input, more-state-than-input) is covered twice: + // * once against a real UDAF, to round-trip an actual aggregate end + // to end through `Partial -> PartialReduce -> Final`; and + // * once against [`InputTypeAssertingUdaf`], whose input / state / + // output types are deliberately pairwise-disjoint within each test + // so a regression that swapped state-field types for input-field + // types (or vice versa) fails the assertion instead of slipping + // through on a coincidental type match. + // + // The stub variants do the heavy lifting on the contract; the real + // ones make sure no real aggregate is broken by it. + // ------------------------------------------------------------------- + + /// Real-UDAF round-trip: aggregate with a single state field and a + /// single input argument (`SUM(b)` — state and input are both `Float64`). + #[tokio::test] + async fn test_partial_reduce_with_single_state_field_and_single_input_arg() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + Ok(vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("SUM(b)") + .build()?, + )]) + }) + .await?; + + // Expected: group 1 -> 10+40=50, group 2 -> 20+50=70, group 3 -> 30+60=90 + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+--------+ + | a | SUM(b) | + +---+--------+ + | 1 | 50.0 | + | 2 | 70.0 | + | 3 | 90.0 | + +---+--------+ + "); + + Ok(()) + } + + /// Real-UDAF round-trip: aggregate with multiple state fields and a + /// single input argument (`AVG(b)` — state is `[sum: Float64, count: + /// UInt64]`). + #[tokio::test] + async fn test_partial_reduce_with_multiple_state_fields_and_single_input_arg() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + Ok(vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("AVG(b)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+--------+ + | a | AVG(b) | + +---+--------+ + | 1 | 25.0 | + | 2 | 35.0 | + | 3 | 45.0 | + +---+--------+ + "); + + Ok(()) + } + + /// Real-UDAF round-trip: aggregate whose state has more fields than the + /// input has arguments (`approx_percentile_cont` carries a t-digest). + #[tokio::test] + async fn test_partial_reduce_with_more_state_fields_than_input_args() -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + Ok(vec![Arc::new( + AggregateExprBuilder::new( + approx_percentile_cont_udaf(), + vec![col("b", schema)?, lit(0.75f32)], + ) + .schema(Arc::clone(schema)) + .alias("approx_percentile_cont(b, 0.75)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+---------------------------------+ + | a | approx_percentile_cont(b, 0.75) | + +---+---------------------------------+ + | 1 | 40.0 | + | 2 | 50.0 | + | 3 | 60.0 | + +---+---------------------------------+ + "); + + Ok(()) + } + + /// Stub variant of + /// [`test_partial_reduce_with_single_state_field_and_single_input_arg`] + /// with disjoint input / state / output types. + /// + /// - input: `Float64` + /// - state: `Int32` + /// - output: `Int64` + /// + /// Any mode that accidentally forwarded state-field types in place of + /// input-field types would fail the assertion in + /// [`InputTypeAssertingUdaf`] instead of being masked by a coincidental + /// type match. + #[tokio::test] + async fn test_partial_reduce_with_single_state_field_and_single_input_arg_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64], + vec![DataType::Int32], + DataType::Int64, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b)") + .build()?, + )]) + }) + .await?; + + // Pipeline completing without error is the real assertion. The + // snapshot guards against silent regressions in the row shape. + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+-------------------------+ + | a | input_type_asserting(b) | + +---+-------------------------+ + | 1 | 0 | + | 2 | 0 | + | 3 | 0 | + +---+-------------------------+ + "); + + Ok(()) + } + + /// Stub variant of + /// [`test_partial_reduce_with_multiple_state_fields_and_single_input_arg`] + /// with disjoint input / state / output types. + /// + /// - input: `Float64` + /// - state: `[Int32, Utf8]` + /// - output: `Int64` + #[tokio::test] + async fn test_partial_reduce_with_multiple_state_fields_and_single_input_arg_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64], + vec![DataType::Int32, DataType::Utf8], + DataType::Int64, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+-------------------------+ + | a | input_type_asserting(b) | + +---+-------------------------+ + | 1 | 0 | + | 2 | 0 | + | 3 | 0 | + +---+-------------------------+ + "); + + Ok(()) + } + + /// Stub variant of + /// [`test_partial_reduce_with_more_state_fields_than_input_args`] with + /// disjoint input / state / output types — and with multiple input + /// arguments to exercise the multi-arg path explicitly. + /// + /// - input: `[Float64, Date32]` + /// - state: `[Int32, Utf8, Boolean]` + /// - output: `Int64` + #[tokio::test] + async fn test_partial_reduce_with_more_state_fields_than_input_args_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64, DataType::Date32], + vec![DataType::Int32, DataType::Utf8, DataType::Boolean], + DataType::Int64, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new( + udaf, + vec![col("b", schema)?, lit(ScalarValue::Date32(Some(1)))], + ) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b, lit)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+------------------------------+ + | a | input_type_asserting(b, lit) | + +---+------------------------------+ + | 1 | 0 | + | 2 | 0 | + | 3 | 0 | + +---+------------------------------+ + "); + + Ok(()) + } + + /// Stub test: many input args, few state fields (5 inputs / 2 state). + /// + /// All eight types involved are pairwise-disjoint: + /// - input: `[Float64, Date32, UInt16, Boolean, Int32]` + /// - state: `[Utf8, Int64]` + /// - output: `Float32` + #[tokio::test] + async fn test_partial_reduce_with_5_input_args_and_2_state_fields_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![ + DataType::Float64, + DataType::Date32, + DataType::UInt16, + DataType::Boolean, + DataType::Int32, + ], + vec![DataType::Utf8, DataType::Int64], + DataType::Float32, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new( + udaf, + vec![ + col("b", schema)?, + lit(ScalarValue::Date32(Some(1))), + lit(ScalarValue::UInt16(Some(1))), + lit(ScalarValue::Boolean(Some(false))), + lit(ScalarValue::Int32(Some(1))), + ], + ) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b, l1, l2, l3, l4)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+-----------------------------------------+ + | a | input_type_asserting(b, l1, l2, l3, l4) | + +---+-----------------------------------------+ + | 1 | 0.0 | + | 2 | 0.0 | + | 3 | 0.0 | + +---+-----------------------------------------+ + "); + + Ok(()) + } + + /// Stub test: few input args, many state fields (2 inputs / 5 state). + /// + /// All eight types involved are pairwise-disjoint: + /// - input: `[Float64, Date32]` + /// - state: `[Boolean, Int32, Utf8, Int64, UInt16]` + /// - output: `Float32` + #[tokio::test] + async fn test_partial_reduce_with_2_input_args_and_5_state_fields_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64, DataType::Date32], + vec![ + DataType::Boolean, + DataType::Int32, + DataType::Utf8, + DataType::Int64, + DataType::UInt16, + ], + DataType::Float32, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new( + udaf, + vec![col("b", schema)?, lit(ScalarValue::Date32(Some(1)))], + ) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b, lit)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+------------------------------+ + | a | input_type_asserting(b, lit) | + +---+------------------------------+ + | 1 | 0.0 | + | 2 | 0.0 | + | 3 | 0.0 | + +---+------------------------------+ + "); + + Ok(()) + } + + /// Test-only aggregate whose `return_type`, `state_fields`, and + /// `accumulator` hooks all assert that they receive the originally- + /// declared input types; the companion accumulator further asserts + /// `update_batch` sees inputs and `merge_batch` sees state. + /// + /// Each test instantiates it with input / state / output types that + /// are pairwise-disjoint, so a regression that forwarded the wrong + /// types fails on type mismatch rather than passing by accident. + #[derive(Debug, PartialEq, Eq, Hash)] + struct InputTypeAssertingUdaf { + signature: Signature, + input_types: Vec<DataType>, + state_types: Vec<DataType>, + output_type: DataType, + } + + fn assert_data_types( + what: &str, + expected: &[DataType], + actual: &[DataType], + ) -> Result<()> { + if actual != expected { + return internal_err!( + "InputTypeAssertingUdaf: {} expected types {:?} but got {:?} — a regression is leaking the wrong types into the accumulator contract", + what, + expected, + actual + ); + } + Ok(()) + } + + /// Produce a zeroed [`ScalarValue`] for `dt`. Only the data types the + /// tests above plug into [`InputTypeAssertingUdaf`] are listed; adding + /// a new type to a test requires extending this match. + fn zero_scalar_for(dt: &DataType) -> Result<ScalarValue> { + match dt { + DataType::Boolean => Ok(ScalarValue::Boolean(Some(false))), + DataType::Int32 => Ok(ScalarValue::Int32(Some(0))), + DataType::Int64 => Ok(ScalarValue::Int64(Some(0))), + DataType::UInt16 => Ok(ScalarValue::UInt16(Some(0))), + DataType::Float32 => Ok(ScalarValue::Float32(Some(0.0))), + DataType::Utf8 => Ok(ScalarValue::Utf8(Some(String::new()))), + other => internal_err!( + "InputTypeAssertingUdaf: no zero ScalarValue registered for {other:?} \ + — extend `zero_scalar_for` when adding a new state/output type" + ), + } + } + + impl InputTypeAssertingUdaf { + fn new( + input_types: Vec<DataType>, + state_types: Vec<DataType>, + output_type: DataType, + ) -> Self { + // Within-test type-disjointness is enforced by construction so + // a future test author can't quietly reintroduce overlap. + assert!( + all_pairwise_distinct(&input_types, &state_types, &output_type), + "InputTypeAssertingUdaf::new: input ({input_types:?}), state \ + ({state_types:?}), and output ({output_type:?}) types must be \ + pairwise-disjoint to avoid accidental passes", + ); + Self { + signature: Signature::exact(input_types.clone(), Volatility::Immutable), + input_types, + state_types, + output_type, + } + } + } + + /// True iff every type in `inputs ∪ states ∪ {output}` is unique. + fn all_pairwise_distinct( + inputs: &[DataType], + states: &[DataType], + output: &DataType, + ) -> bool { + let mut seen = HashSet::new(); + for dt in inputs + .iter() + .chain(states.iter()) + .chain(std::iter::once(output)) + { + if !seen.insert(dt) { + return false; + } + } + true + } + + impl AggregateUDFImpl for InputTypeAssertingUdaf { + fn name(&self) -> &str { + "input_type_asserting" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> { + assert_data_types("return_type(arg_types)", &self.input_types, arg_types)?; + Ok(self.output_type.clone()) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + let actual: Vec<DataType> = args + .input_fields + .iter() + .map(|f| f.data_type().clone()) + .collect(); + assert_data_types( + "state_fields(args.input_fields)", + &self.input_types, + &actual, + )?; + Ok(self + .state_types + .iter() + .enumerate() + .map(|(i, dt)| { + Field::new(format!("{}[s{i}]", args.name), dt.clone(), true).into() + }) + .collect()) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + let actual: Vec<DataType> = acc_args + .expr_fields + .iter() + .map(|f| f.data_type().clone()) + .collect(); + assert_data_types( + "accumulator(acc_args.expr_fields)", + &self.input_types, + &actual, + )?; + Ok(Box::new(InputTypeAssertingAccumulator { + input_types: self.input_types.clone(), + state_types: self.state_types.clone(), + output_type: self.output_type.clone(), + })) + } + } + + /// Companion accumulator for [`InputTypeAssertingUdaf`]. + /// + /// - `update_batch` must always receive arrays of the original input + /// types. + /// - `merge_batch` must always receive arrays of the declared state + /// types. + /// + /// Anything else means a non-input mode is calling the wrong path. + #[derive(Debug)] + struct InputTypeAssertingAccumulator { + input_types: Vec<DataType>, + state_types: Vec<DataType>, + output_type: DataType, + } + + impl Accumulator for InputTypeAssertingAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let actual: Vec<DataType> = + values.iter().map(|a| a.data_type().clone()).collect(); + assert_data_types("update_batch(values)", &self.input_types, &actual) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + zero_scalar_for(&self.output_type) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + self.state_types.iter().map(zero_scalar_for).collect() + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let actual: Vec<DataType> = + states.iter().map(|a| a.data_type().clone()).collect(); + assert_data_types("merge_batch(states)", &self.state_types, &actual) + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct NoFirstEmitUdaf { + signature: Signature, + } + + impl NoFirstEmitUdaf { + fn new() -> Self { + Self { + signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable), + } + } + } + + impl AggregateUDFImpl for NoFirstEmitUdaf { + fn name(&self) -> &str { + "no_first_emit" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { + Ok(DataType::Int64) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + Ok(vec![Arc::new(Field::new( + format!("{}[count]", args.name), + DataType::Int64, + false, + ))]) + } + + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(NoFirstEmitAccumulator)) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + _args: AccumulatorArgs, + ) -> Result<Box<dyn GroupsAccumulator>> { + Ok(Box::new(NoFirstEmitGroupsAccumulator { counts: vec![] })) + } + } + + #[derive(Debug)] + struct NoFirstEmitAccumulator; + + impl Accumulator for NoFirstEmitAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + Ok(ScalarValue::Int64(Some(0))) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result<Vec<ScalarValue>> { + Ok(vec![ScalarValue::Int64(Some(0))]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[derive(Debug)] + struct NoFirstEmitGroupsAccumulator { + counts: Vec<i64>, + } + + impl NoFirstEmitGroupsAccumulator { + fn emit_counts(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + match emit_to { + EmitTo::All => { + let counts = std::mem::take(&mut self.counts); + Ok(Arc::new(Int64Array::from(counts))) + } + EmitTo::First(_) => internal_err!( + "partial grouped aggregate output must materialize with EmitTo::All before slicing" + ), + } + } + } + + impl GroupsAccumulator for NoFirstEmitGroupsAccumulator { + fn update_batch( + &mut self, + _values: &[ArrayRef], + group_indices: &[usize], + _opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.counts.resize(total_num_groups, 0); + for group_index in group_indices { + self.counts[*group_index] += 1; + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> { + self.emit_counts(emit_to) + } + + fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> { + Ok(vec![self.emit_counts(emit_to)?]) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result<Vec<ArrayRef>> { + assert_eq!(values.len(), 1, "one argument to convert_to_state"); + let counts = match opt_filter { + Some(filter) => filter + .iter() + .map(|value| i64::from(value.unwrap_or(false))) + .collect::<Vec<_>>(), + None => vec![1; values[0].len()], + }; + Ok(vec![Arc::new(Int64Array::from(counts))]) + } + + fn merge_batch( + &mut self, + _values: &[ArrayRef], + _group_indices: &[usize], + _total_num_groups: usize, + ) -> Result<()> { + Ok(()) + } + + fn size(&self) -> usize { + size_of_val(self) + self.counts.capacity() * size_of::<i64>() + } + } + + /// Test that [`AggregateExec::with_dynamic_filter_expr`] overrides the existing dynamic filter + #[test] + fn test_with_dynamic_filter() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + // Partial min aggregate supports dynamic filtering + let agg = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![]), + vec![Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("min_a") + .build()?, + )], + vec![None], + child, + Arc::clone(&schema), + )?; + + // Assertion 1: A filter with the same children can override the existing + // dynamic filter. + let new_df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![col("a", &schema)?], + lit(false), + )); + let agg = agg.with_dynamic_filter_expr(Arc::clone(&new_df))?; + let produced = agg.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), new_df.expression_id()); + + // The aggregate's filter should now resolve to the new inner expression. + let swapped = produced[0] + .downcast_ref::<DynamicFilterPhysicalExpr>() + .expect("produced expression should be a DynamicFilterPhysicalExpr") + .current()?; + assert_eq!(format!("{swapped}"), format!("{}", lit(false))); + + // Assertion 2: A filter that has been through `PhysicalExpr::with_new_children` + // should still be accepted when the new children are equivalent to the originals. + let new_df_as_pexpr: Arc<dyn PhysicalExpr> = + Arc::<DynamicFilterPhysicalExpr>::clone(&new_df); + let remapped_pexpr = + new_df_as_pexpr.with_new_children(vec![col("a", &schema)?])?; + let Ok(remapped_df) = (remapped_pexpr as Arc<dyn std::any::Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + else { + panic!("should be DynamicFilterPhysicalExpr after with_new_children"); + }; + // Hard to assert this because the filter is identical. No error means + // the filter was accepted. That's a good enough assertion for now. + let _agg = agg.with_dynamic_filter_expr(remapped_df)?; + Ok(()) + } + + #[test] + fn test_plan_contains_expression_id_recurses_plans_and_expressions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let empty: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![col("a", &schema)?], + lit(true), + )); + let expression_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + + assert!(!plan_contains_expression_id(&empty, expression_id)?); + + let dynamic_filter_expr: Arc<dyn PhysicalExpr> = + Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter); + let predicate: Arc<dyn PhysicalExpr> = + Arc::new(NotExpr::new(dynamic_filter_expr)); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExecBuilder::new(predicate, empty).build()?); + let projection: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new( + [ProjectionExpr::new_from_expression( + col("a", &schema)?, + &schema, + )?], + filter, + )?); + + assert!(plan_contains_expression_id(&projection, expression_id)?); + Ok(()) + } + + /// Test that [`AggregateExec::with_dynamic_filter_expr`] errors when the aggregate does not support dynamic filtering + #[test] + fn test_with_dynamic_filter_error_unsupported() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + // Final mode with a group-by does not support dynamic filters. + let agg = AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]), + vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("sum_b") + .build()?, + )], + vec![None], + child, + Arc::clone(&schema), + )?; + assert!(agg.dynamic_expressions_produced().is_empty()); + + let df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![col("a", &schema)?], + lit(true), + )); + assert!(agg.with_dynamic_filter_expr(df).is_err()); + Ok(()) + } + + /// Test that [`AggregateExec::with_dynamic_filter_expr`] errors when the column is not in the schema + #[test] + fn test_with_dynamic_filter_error_column_mismatch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let agg = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![]), + vec![Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("min_a") + .build()?, + )], + vec![None], + child, + Arc::clone(&schema), + )?; + + let df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("bad", 99)) as _], + lit(true), + )); + assert!(agg.with_dynamic_filter_expr(df).is_err()); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/order/full.rs b/native/vendor/datafusion-physical-plan/src/aggregates/order/full.rs new file mode 100644 index 0000000..ca818d6 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/order/full.rs
@@ -0,0 +1,156 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion_expr::EmitTo; +use std::mem::size_of; + +/// Tracks grouping state when the data is ordered entirely by its +/// group keys +/// +/// When the group values are sorted, as soon as we see group `n+1` we +/// know we will never see any rows for group `n` again and thus they +/// can be emitted. +/// +/// For example, given `SUM(amt) GROUP BY id` if the input is sorted +/// by `id` as soon as a new `id` value is seen all previous values +/// can be emitted. +/// +/// The state is tracked like this: +/// +/// ```text +/// ┌─────┐ ┌──────────────────┐ +/// │┌───┐│ │ ┌──────────────┐ │ ┏━━━━━━━━━━━━━━┓ +/// ││ 0 ││ │ │ 123 │ │ ┌─────┃ 13 ┃ +/// │└───┘│ │ └──────────────┘ │ │ ┗━━━━━━━━━━━━━━┛ +/// │ ... │ │ ... │ │ +/// │┌───┐│ │ ┌──────────────┐ │ │ current +/// ││12 ││ │ │ 234 │ │ │ +/// │├───┤│ │ ├──────────────┤ │ │ +/// ││12 ││ │ │ 234 │ │ │ +/// │├───┤│ │ ├──────────────┤ │ │ +/// ││13 ││ │ │ 456 │◀┼───┘ +/// │└───┘│ │ └──────────────┘ │ +/// └─────┘ └──────────────────┘ +/// +/// group indices group_values current tracks the most +/// (in group value recent group index +/// order) +/// ``` +/// +/// In this diagram, the current group is `13`, and thus groups +/// `0..12` can be emitted. Note that `13` can not yet be emitted as +/// there may be more values in the next batch with the same group_id. +#[derive(Debug)] +pub struct GroupOrderingFull { + state: State, +} + +#[derive(Debug)] +enum State { + /// Seen no input yet + Start, + + /// Data is in progress. `current` is the current group for which + /// values are being generated. Can emit `current` - 1 + InProgress { current: usize }, + + /// Seen end of input: all groups can be emitted + Complete, +} + +impl GroupOrderingFull { + pub fn new() -> Self { + Self { + state: State::Start, + } + } + + // How many groups be emitted, or None if no data can be emitted + pub fn emit_to(&self) -> Option<EmitTo> { + match &self.state { + State::Start => None, + State::InProgress { current, .. } => { + if *current == 0 { + // Can not emit if still on the first row + None + } else { + // otherwise emit all rows prior to the current group + Some(EmitTo::First(*current)) + } + } + State::Complete => Some(EmitTo::All), + } + } + + /// remove the first n groups from the internal state, shifting + /// all existing indexes down by `n` + pub fn remove_groups(&mut self, n: usize) { + match &mut self.state { + State::Start => panic!("invalid state: start"), + State::InProgress { current } => { + // shift down by n + assert!(*current >= n); + *current -= n; + } + State::Complete => panic!("invalid state: complete"), + } + } + + /// Note that the input is complete so any outstanding groups are done as well + pub fn input_done(&mut self) { + self.state = State::Complete; + } + + /// Starts tracking a new fully ordered input segment. + pub fn reset(&mut self) { + self.state = State::Start; + } + + /// Called when new groups are added in a batch. See documentation + /// on [`super::GroupOrdering::new_groups`] + pub fn new_groups(&mut self, total_num_groups: usize) { + assert_ne!(total_num_groups, 0); + + // Update state + let max_group_index = total_num_groups - 1; + self.state = match self.state { + State::Start => State::InProgress { + current: max_group_index, + }, + State::InProgress { current } => { + // expect to see new group indexes when called again + assert!(current <= max_group_index, "{current} <= {max_group_index}"); + State::InProgress { + current: max_group_index, + } + } + State::Complete => { + panic!("Saw new group after input was complete"); + } + }; + } + + pub(crate) fn size(&self) -> usize { + size_of::<Self>() + } +} + +impl Default for GroupOrderingFull { + fn default() -> Self { + Self::new() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/order/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/order/mod.rs new file mode 100644 index 0000000..259411b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/order/mod.rs
@@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::mem::size_of; + +use arrow::array::ArrayRef; +use datafusion_common::Result; +use datafusion_expr::EmitTo; + +mod full; +mod partial; + +use crate::InputOrderMode; +pub use full::GroupOrderingFull; +pub use partial::GroupOrderingPartial; + +/// Ordering information for each group in the hash table +#[derive(Debug)] +pub enum GroupOrdering { + /// Groups are not ordered + None, + /// Groups are ordered by some pre-set of the group keys + Partial(GroupOrderingPartial), + /// Groups are entirely contiguous, + Full(GroupOrderingFull), +} + +impl GroupOrdering { + /// Create a `GroupOrdering` for the specified ordering + pub fn try_new(mode: &InputOrderMode) -> Result<Self> { + match mode { + InputOrderMode::Linear => Ok(GroupOrdering::None), + InputOrderMode::PartiallySorted(order_indices) => { + GroupOrderingPartial::try_new(order_indices.clone()) + .map(GroupOrdering::Partial) + } + InputOrderMode::Sorted => Ok(GroupOrdering::Full(GroupOrderingFull::new())), + } + } + + /// Returns how many groups can be emitted while respecting the current + /// ordering guarantees, or `None` if no data can be emitted. + pub fn emit_to(&self) -> Option<EmitTo> { + match self { + GroupOrdering::None => None, + GroupOrdering::Partial(partial) => partial.emit_to(), + GroupOrdering::Full(full) => full.emit_to(), + } + } + + /// Returns the emit strategy to use under memory pressure (OOM). + /// + /// Returns the strategy that must be used when emitting up to `n` groups + /// while respecting the current ordering guarantees. + /// + /// Returns `None` if no data can be emitted. + pub fn oom_emit_to(&self, n: usize) -> Option<EmitTo> { + if n == 0 { + return None; + } + + match self { + GroupOrdering::None => Some(EmitTo::First(n)), + GroupOrdering::Partial(_) | GroupOrdering::Full(_) => { + self.emit_to().map(|emit_to| match emit_to { + EmitTo::First(max) => EmitTo::First(n.min(max)), + EmitTo::All => EmitTo::First(n), + }) + } + } + } + + /// Updates the state to indicate that the input is complete. + pub fn input_done(&mut self) { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => partial.input_done(), + GroupOrdering::Full(full) => full.input_done(), + } + } + + /// Resets the ordering state while preserving the configured ordering mode. + /// + /// Ordered partial aggregation uses this after passing intermediate states + /// downstream, and ordered final aggregation uses it after spilling a run. + /// In both cases the hash table is empty and can start tracking the next + /// input batch from a fresh ordering state. + pub fn reset(&mut self) { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => partial.reset(), + GroupOrdering::Full(full) => full.reset(), + } + } + + /// Removes the first `n` groups from the internal state, shifting all + /// existing indexes down by `n`. + pub fn remove_groups(&mut self, n: usize) { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => partial.remove_groups(n), + GroupOrdering::Full(full) => full.remove_groups(n), + } + } + + /// Called when new groups are added in a batch. + /// + /// * `batch_group_values`: group key values for each row in the batch + /// + /// * `group_indices`: indices for each row in the batch + /// + /// * `total_num_groups`: total number of groups (so max + /// group_index is total_num_groups - 1). + pub fn new_groups( + &mut self, + batch_group_values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => { + partial.new_groups( + batch_group_values, + group_indices, + total_num_groups, + )?; + } + GroupOrdering::Full(full) => { + full.new_groups(total_num_groups); + } + }; + Ok(()) + } + + /// Returns the size of memory used by the ordering state, in bytes. + pub fn size(&self) -> usize { + size_of::<Self>() + + match self { + GroupOrdering::None => 0, + GroupOrdering::Partial(partial) => partial.size(), + GroupOrdering::Full(full) => full.size(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::sync::Arc; + + use arrow::array::Int32Array; + + #[test] + fn test_oom_emit_to_none_ordering() { + let group_ordering = GroupOrdering::None; + + assert_eq!(group_ordering.oom_emit_to(0), None); + assert_eq!(group_ordering.oom_emit_to(5), Some(EmitTo::First(5))); + } + + /// Creates a partially ordered grouping state with three groups. + /// + /// `sort_key_values` controls whether a sort boundary exists in the batch: + /// distinct values such as `[1, 2, 3]` create boundaries, while repeated + /// values such as `[1, 1, 1]` do not. + fn partial_ordering(sort_key_values: Vec<i32>) -> Result<GroupOrdering> { + let mut group_ordering = + GroupOrdering::Partial(GroupOrderingPartial::try_new(vec![0])?); + + let batch_group_values: Vec<ArrayRef> = vec![ + Arc::new(Int32Array::from(sort_key_values)), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ]; + let group_indices = vec![0, 1, 2]; + + group_ordering.new_groups(&batch_group_values, &group_indices, 3)?; + + Ok(group_ordering) + } + + #[test] + fn test_oom_emit_to_partial_clamps_to_boundary() -> Result<()> { + let group_ordering = partial_ordering(vec![1, 2, 3])?; + + // Can emit both `1` and `2` groups because we have seen `3` + assert_eq!(group_ordering.emit_to(), Some(EmitTo::First(2))); + assert_eq!(group_ordering.oom_emit_to(1), Some(EmitTo::First(1))); + assert_eq!(group_ordering.oom_emit_to(3), Some(EmitTo::First(2))); + + Ok(()) + } + + #[test] + fn test_oom_emit_to_partial_without_boundary() -> Result<()> { + let group_ordering = partial_ordering(vec![1, 1, 1])?; + + // Can't emit the last `1` group as it may have more values + assert_eq!(group_ordering.emit_to(), None); + assert_eq!(group_ordering.oom_emit_to(3), None); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/order/partial.rs b/native/vendor/datafusion-physical-plan/src/aggregates/order/partial.rs new file mode 100644 index 0000000..1603bb6 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/order/partial.rs
@@ -0,0 +1,358 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cmp::Ordering; +use std::mem::size_of; +use std::sync::Arc; + +use arrow::array::ArrayRef; +use arrow::compute::SortOptions; +use arrow_ord::partition::partition; +use datafusion_common::utils::{compare_rows, get_row_at_idx}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +/// Tracks grouping state when the data is ordered by some subset of +/// the group keys. +/// +/// Once the next *sort key* value is seen, never see groups with that +/// sort key again, so we can emit all groups with the previous sort +/// key and earlier. +/// +/// For example, given `SUM(amt) GROUP BY id, state` if the input is +/// sorted by `state`, when a new value of `state` is seen, all groups +/// with prior values of `state` can be emitted. +/// +/// The state is tracked like this: +/// +/// ```text +/// ┏━━━━━━━━━━━━━━━━━┓ ┏━━━━━━━┓ +/// ┌─────┐ ┌───────────────────┐ ┌─────┃ 9 ┃ ┃ "MD" ┃ +/// │┌───┐│ │ ┌──────────────┐ │ │ ┗━━━━━━━━━━━━━━━━━┛ ┗━━━━━━━┛ +/// ││ 0 ││ │ │ 123, "MA" │ │ │ current_sort sort_key +/// │└───┘│ │ └──────────────┘ │ │ +/// │ ... │ │ ... │ │ current_sort tracks the +/// │┌───┐│ │ ┌──────────────┐ │ │ smallest group index that had +/// ││ 8 ││ │ │ 765, "MA" │ │ │ the same sort_key as current +/// │├───┤│ │ ├──────────────┤ │ │ +/// ││ 9 ││ │ │ 923, "MD" │◀─┼─┘ +/// │├───┤│ │ ├──────────────┤ │ ┏━━━━━━━━━━━━━━┓ +/// ││10 ││ │ │ 345, "MD" │ │ ┌─────┃ 11 ┃ +/// │├───┤│ │ ├──────────────┤ │ │ ┗━━━━━━━━━━━━━━┛ +/// ││11 ││ │ │ 124, "MD" │◀─┼──┘ current +/// │└───┘│ │ └──────────────┘ │ +/// └─────┘ └───────────────────┘ +/// +/// group indices +/// (in group value group_values current tracks the most +/// order) recent group index +/// ``` +#[derive(Debug)] +pub struct GroupOrderingPartial { + /// State machine + state: State, + + /// The indexes of the group by columns that form the sort key. + /// For example if grouping by `id, state` and ordered by `state` + /// this would be `[1]`. + order_indices: Vec<usize>, +} + +#[derive(Debug, Default, PartialEq)] +enum State { + /// The ordering was temporarily taken. `Self::Taken` is left + /// when state must be temporarily taken to satisfy the borrow + /// checker. If an error happens before the state can be restored, + /// the ordering information is lost and execution can not + /// proceed, but there is no undefined behavior. + #[default] + Taken, + + /// Seen no input yet + Start, + + /// Data is in progress. + InProgress { + /// Smallest group index with the sort_key + current_sort: usize, + /// The sort key of group_index `current_sort` + sort_key: Vec<ScalarValue>, + /// index of the current group for which values are being + /// generated + current: usize, + }, + + /// Seen end of input, all groups can be emitted + Complete, +} + +impl State { + fn size(&self) -> usize { + match self { + State::Taken => 0, + State::Start => 0, + State::InProgress { sort_key, .. } => sort_key + .iter() + .map(|scalar_value| scalar_value.size()) + .sum(), + State::Complete => 0, + } + } +} + +impl GroupOrderingPartial { + /// TODO: Remove unnecessary `input_schema` parameter. + pub fn try_new(order_indices: Vec<usize>) -> Result<Self> { + debug_assert!(!order_indices.is_empty()); + Ok(Self { + state: State::Start, + order_indices, + }) + } + + /// Select sort keys from the group values + /// + /// For example, if group_values had `A, B, C` but the input was + /// only sorted on `B` and `C` this should return rows for (`B`, + /// `C`) + fn compute_sort_keys(&mut self, group_values: &[ArrayRef]) -> Vec<ArrayRef> { + // Take only the columns that are in the sort key + self.order_indices + .iter() + .map(|&idx| Arc::clone(&group_values[idx])) + .collect() + } + + /// How many groups be emitted, or None if no data can be emitted + pub fn emit_to(&self) -> Option<EmitTo> { + match &self.state { + State::Taken => unreachable!("State previously taken"), + State::Start => None, + State::InProgress { current_sort, .. } => { + // Can not emit if we are still on the first row sort + // row otherwise we can emit all groups that had earlier sort keys + // + if *current_sort == 0 { + None + } else { + Some(EmitTo::First(*current_sort)) + } + } + State::Complete => Some(EmitTo::All), + } + } + + /// remove the first n groups from the internal state, shifting + /// all existing indexes down by `n` + pub fn remove_groups(&mut self, n: usize) { + match &mut self.state { + State::Taken => unreachable!("State previously taken"), + State::Start => panic!("invalid state: start"), + State::InProgress { + current_sort, + current, + sort_key: _, + } => { + // shift indexes down by n + assert!(*current >= n); + *current -= n; + assert!(*current_sort >= n); + *current_sort -= n; + } + State::Complete => panic!("invalid state: complete"), + } + } + + /// Note that the input is complete so any outstanding groups are done as well + pub fn input_done(&mut self) { + self.state = match self.state { + State::Taken => unreachable!("State previously taken"), + _ => State::Complete, + }; + } + + /// Starts tracking a new ordered input segment with the same sort-key + /// columns. + pub fn reset(&mut self) { + self.state = State::Start; + } + + fn updated_sort_key( + current_sort: usize, + sort_key: Option<Vec<ScalarValue>>, + range_current_sort: usize, + range_sort_key: Vec<ScalarValue>, + ) -> Result<(usize, Vec<ScalarValue>)> { + if let Some(sort_key) = sort_key { + let sort_options = vec![SortOptions::new(false, false); sort_key.len()]; + let ordering = compare_rows(&sort_key, &range_sort_key, &sort_options)?; + if ordering == Ordering::Equal { + return Ok((current_sort, sort_key)); + } + } + + Ok((range_current_sort, range_sort_key)) + } + + /// Called when new groups are added in a batch. See documentation + /// on [`super::GroupOrdering::new_groups`] + pub fn new_groups( + &mut self, + batch_group_values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + assert!(total_num_groups > 0); + assert!(!batch_group_values.is_empty()); + + let max_group_index = total_num_groups - 1; + + let (current_sort, sort_key) = match std::mem::take(&mut self.state) { + State::Taken => unreachable!("State previously taken"), + State::Start => (0, None), + State::InProgress { + current_sort, + sort_key, + .. + } => (current_sort, Some(sort_key)), + State::Complete => { + panic!("Saw new group after the end of input"); + } + }; + + // Select the sort key columns + let sort_keys = self.compute_sort_keys(batch_group_values); + + // Check if the sort keys indicate a boundary inside the batch + let ranges = partition(&sort_keys)?.ranges(); + let last_range = ranges.last().unwrap(); + + let range_current_sort = group_indices[last_range.start]; + let range_sort_key = get_row_at_idx(&sort_keys, last_range.start)?; + + let (current_sort, sort_key) = if last_range.start == 0 { + // There was no boundary in the batch. Compare with the previous sort_key (if present) + // to check if there was a boundary between the current batch and the previous one. + Self::updated_sort_key( + current_sort, + sort_key, + range_current_sort, + range_sort_key, + )? + } else { + (range_current_sort, range_sort_key) + }; + + self.state = State::InProgress { + current_sort, + current: max_group_index, + sort_key, + }; + + Ok(()) + } + + /// Return the size of memory allocated by this structure + pub(crate) fn size(&self) -> usize { + size_of::<Self>() + self.order_indices.allocated_size() + self.state.size() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::Int32Array; + + #[test] + fn test_group_ordering_partial() -> Result<()> { + // Ordered on column a + let order_indices = vec![0]; + let mut group_ordering = GroupOrderingPartial::try_new(order_indices)?; + + let batch_group_values: Vec<ArrayRef> = vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![2, 1, 3])), + ]; + + let group_indices = vec![0, 1, 2]; + let total_num_groups = 3; + + group_ordering.new_groups( + &batch_group_values, + &group_indices, + total_num_groups, + )?; + + assert_eq!( + group_ordering.state, + State::InProgress { + current_sort: 2, + sort_key: vec![ScalarValue::Int32(Some(3))], + current: 2 + } + ); + + // push without a boundary + let batch_group_values: Vec<ArrayRef> = vec![ + Arc::new(Int32Array::from(vec![3, 3, 3])), + Arc::new(Int32Array::from(vec![2, 1, 7])), + ]; + let group_indices = vec![3, 4, 5]; + let total_num_groups = 6; + + group_ordering.new_groups( + &batch_group_values, + &group_indices, + total_num_groups, + )?; + + assert_eq!( + group_ordering.state, + State::InProgress { + current_sort: 2, + sort_key: vec![ScalarValue::Int32(Some(3))], + current: 5 + } + ); + + // push with only a boundary to previous batch + let batch_group_values: Vec<ArrayRef> = vec![ + Arc::new(Int32Array::from(vec![4, 4, 4])), + Arc::new(Int32Array::from(vec![1, 1, 1])), + ]; + let group_indices = vec![6, 7, 8]; + let total_num_groups = 9; + + group_ordering.new_groups( + &batch_group_values, + &group_indices, + total_num_groups, + )?; + assert_eq!( + group_ordering.state, + State::InProgress { + current_sort: 6, + sort_key: vec![ScalarValue::Int32(Some(4))], + current: 8 + } + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs new file mode 100644 index 0000000..19deedc --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_final_stream.rs
@@ -0,0 +1,902 @@ +// 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. + +//! Final aggregate stream for ordered partial-state input. + +use std::ops::ControlFlow; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable}; +use super::group_values::GroupByMetrics; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Final aggregate stream for `InputOrderMode::Sorted` and +/// `InputOrderMode::PartiallySorted`. +/// +/// See comments at [`super::ordered_partial_stream::OrderedPartialAggregateStream`] for details. +/// +/// # Spilling +/// +/// This section is only for implementation notes, for background, see [`super::ordered_partial_stream::OrderedPartialAggregateStream`] +/// +/// For partially sorted input, spilling works as follows: +/// +/// - Reserve the table footprint plus one `u32` sort index per buffered group. The +/// extra index array is used in later sorting before spilling. +/// - On memory pressure, materialize all group states into one batch. +/// - Use [`IncrementalSortIterator`] to compute the full-batch index, then +/// materialize and write one sorted `batch_size` slice at a time. The original +/// batch and full index remain live until the run is written. +/// - After input ends, merge the sorted runs and replay them through a fully +/// ordered final aggregate stream. +pub(crate) struct OrderedFinalAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + state: Option<OrderedFinalAggregateState>, +} + +/// Spill configuration and accumulated runs for partially ordered final +/// aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct OrderedFinalSpillContext { + /// Aggregate configuration + agg: AggregateExec, + /// Task context + context: Arc<TaskContext>, + /// Original partition index + partition: usize, + /// Target batch size from configuration + batch_size: usize, + /// Full group-key ordering, such ordering with be kept in: a) individual spill + /// files, b) order after final merging and streaming aggregate + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Fully sorted spill runs waiting to be merged. + spills: Vec<SortedSpillFile>, +} + +/// See comments at `poll_next()` for details. +enum OrderedFinalAggregateState { + ReadingInput { + table: OrderedAggregateTable<FinalMarker>, + /// None if either + /// - Disk Manager doesn't enable temporary file creation + /// - The group keys are fully ordered, it's expected to use bounded memory + spill_context: Option<Box<OrderedFinalSpillContext>>, + }, + Spilling { + table: OrderedAggregateTable<FinalMarker>, + spill_context: Box<OrderedFinalSpillContext>, + }, + ProducingOutput { + table: OrderedAggregateTable<FinalMarker>, + }, + PreparingMergeInput { + table: OrderedAggregateTable<FinalMarker>, + spill_context: Box<OrderedFinalSpillContext>, + }, + MergingSpills { + stream: SendableRecordBatchStream, + }, + Done, +} + +type OrderedFinalAggregatePoll = Poll<Option<Result<RecordBatch>>>; +type OrderedFinalAggregateStateTransition = ControlFlow< + (OrderedFinalAggregatePoll, OrderedFinalAggregateState), + OrderedFinalAggregateState, +>; + +impl OrderedFinalSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + batch_size: usize, + input_order_mode: &InputOrderMode, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result<Self> { + let group_schema = agg.group_by.group_schema(spill_schema)?; + let output_ordering = agg.cache.output_ordering(); + let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { + return internal_err!("Ordered final spill requires partially ordered input"); + }; + let spill_indices = order_indices.iter().copied().chain( + (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), + ); + let spill_sort_exprs = spill_indices.map(|idx| { + let field = group_schema.field(idx); + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Ordered final spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + Ok(Self { + agg: agg.clone(), + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`OrderedFinalAggregateStream`] for spilling details. + fn spill_table( + &mut self, + table: &mut OrderedAggregateTable<FinalMarker>, + ) -> Result<()> { + let Some(batch) = table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "OrderedFinalAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Ordered final aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run and finalizes it through the fully ordered path. + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result<SendableRecordBatchStream> { + let Self { + agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl OrderedFinalAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug_assert!(matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + )); + debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + + let input = agg.input.execute(partition, Arc::clone(context))?; + Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) + } + + pub(in crate::aggregates) fn new_with_input( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + input: SendableRecordBatchStream, + input_order_mode: &InputOrderMode, + ) -> Result<Self> { + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reservation = + MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + // HACK: Technically, fully ordered aggregate is a non-spillable + // consumer, since it uses bounded memory. There is a known race + // condition bug, and we set it to spillable to let it have larger + // memory budget to suppress the bug. + // Bug issue: https://github.com/apache/datafusion/issues/17334 + .with_can_spill(true) + .register(context.memory_pool()); + Self::new_with_input_and_metrics( + agg, + context, + partition, + input, + input_order_mode, + baseline_metrics, + group_by_metrics, + Some(spill_metrics), + reservation, + ) + } + + #[expect( + clippy::too_many_arguments, + reason = "keeps replay metric reuse explicit" + )] + /// Builds the stream with the reservation of its logical aggregate operator. + /// Replay callers pass a sibling of the reservation used by the merge input, + /// keeping both components under one memory-consumer registration. + pub(in crate::aggregates) fn new_with_input_and_metrics( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + input: SendableRecordBatchStream, + input_order_mode: &InputOrderMode, + baseline_metrics: BaselineMetrics, + group_by_metrics: GroupByMetrics, + spill_metrics: Option<SpillMetrics>, + reservation: MemoryReservation, + ) -> Result<Self> { + debug_assert!(matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + )); + debug_assert_ne!(*input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + + let can_spill = matches!(input_order_mode, InputOrderMode::PartiallySorted(_)) + && context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + let Some(spill_metrics) = spill_metrics else { + return internal_err!("Spillable ordered final stream requires metrics"); + }; + Some(Box::new(OrderedFinalSpillContext::new( + agg, + context, + partition, + batch_size, + input_order_mode, + &input_schema, + spill_metrics, + )?)) + } else { + None + }; + + let table = OrderedAggregateTable::<FinalMarker>::new_with_input_order( + agg, + &input_schema, + Arc::clone(&schema), + batch_size, + input_order_mode, + group_by_metrics, + )?; + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + state: Some(OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_internal_err(message: &str) -> OrderedFinalAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(internal_err!("{message}"))), + OrderedFinalAggregateState::Done, + )) + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + table: &OrderedAggregateTable<FinalMarker>, + spill_context: Option<&OrderedFinalSpillContext>, + ) -> usize { + let table_size = table.memory_size(); + if spill_context.is_some() { + // See `OrderedFinalAggregateStream` comments for how is it estimated + table_size.saturating_add(table.num_groups().saturating_mul(size_of::<u32>())) + } else { + table_size + } + } + + /// Consumes one ordered partial-state input batch, then immediately emits + /// finalized groups if the ordering proves any group is ready. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::ReadingInput { + mut table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected ReadingInput state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )), + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )); + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )); + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(spill_context) = spill_context else { + // `None` means spilling is not supported, see comments + // at `OrderedFinalAggregateState` for details. + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + }; + if table.is_empty() { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + return ControlFlow::Continue( + OrderedFinalAggregateState::Spilling { + table, + spill_context, + }, + ); + } + Err(e) => { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + } + + let result = if spill_context + .as_ref() + .is_some_and(|spill_context| spill_context.has_spills()) + { + // Once one incomplete run is spilled, every remaining state + // must participate in replay so no group is finalized twice. + Ok(None) + } else { + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + result + }; + + match result { + // Some finalized groups can be emitted. Yield them, then + // continue aggregating input in the current state. + Ok(Some(batch)) => { + if let Err(e) = + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )) + { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + let next_state = OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok( + batch.record_output(&self.baseline_metrics) + ))), + next_state, + )) + } + // Can't do early emit, continue aggregating. + Ok(None) => { + ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )), + } + } + Poll::Ready(Some(Err(e))) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )), + Poll::Ready(None) => { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + OrderedFinalAggregateState::PreparingMergeInput { + table, + spill_context, + }, + ) + } + _ => { + table.input_done(); + ControlFlow::Continue( + OrderedFinalAggregateState::ProducingOutput { table }, + ) + } + } + } + } + } + + /// 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( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::Spilling { + mut table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected Spilling state", + ); + }; + + // Sanity check: it's impossible to OOM when the table is empty + if table.is_empty() { + return ControlFlow::Break(( + Poll::Ready(Some(internal_err!( + "Ordered final aggregation entered Spilling with an empty table" + ))), + OrderedFinalAggregateState::Done, + )); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + result = + Err(e.context("Decreasing allocation after spilling should succeed")); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input + Ok(()) => ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + table, + spill_context: Some(spill_context), + }), + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + } + } + + /// 1. Spills the last in-memory run. + /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// merge to all spills. + /// 3. Constructs a replay stream: an ordered 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( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::PreparingMergeInput { + mut table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered final 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 table) { + Ok(()) => { + let group_by_metrics = table.group_by_metrics(); + drop(table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ), + Err(e) => Err(e), + } + } + Err(e) => Err(e), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(OrderedFinalAggregateState::MergingSpills { + stream, + }) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + } + } + + /// 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: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::MergingSpills { mut stream } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedFinalAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + OrderedFinalAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + Poll::Ready(None) => ControlFlow::Continue(OrderedFinalAggregateState::Done), + } + } + + /// Emits one batch after input is exhausted. + /// + /// `table.input_done()` has already made every remaining group safe to emit, + /// so this state keeps draining until the table is empty. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::ProducingOutput { table } = original_state else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected ProducingOutput state", + ); + }; + + let mut table = table; + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let next_state = if table.is_empty() { + drop(table); + if let Err(e) = self.reservation.try_resize(0) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + OrderedFinalAggregateState::Done + } else { + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ProducingOutput { table }, + )); + } + OrderedFinalAggregateState::ProducingOutput { table } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ProducingOutput { table }, + )), + Ok(None) => { + drop(table); + let next_state = OrderedFinalAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); + } + ControlFlow::Continue(next_state) + } + } + } +} + +impl Stream for OrderedFinalAggregateStream { + type Item = Result<RecordBatch>; + + /// Entry point for the ordered final aggregate state machine. + /// + /// See comments in [`OrderedFinalAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered partial-state input and merging + /// those states into the ordered final aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Merge one input batch. If it fits in memory, optionally yield groups + /// proven complete by the input ordering, then read the next 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. Mark every remaining group as + /// complete and produce its final result. + /// -> 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 remaining final aggregate batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + loop { + let cur_state = self + .state + .take() + .expect("OrderedFinalAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ OrderedFinalAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ OrderedFinalAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ OrderedFinalAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ OrderedFinalAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } + state @ OrderedFinalAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ OrderedFinalAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + // Errors are terminal: discard all operator state and release + // its upstream input and memory reservation before returning. + drop(next_state); + self.close_input(); + self.reservation.free(); + self.state = Some(OrderedFinalAggregateState::Done); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for OrderedFinalAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/ordered_partial_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_partial_stream.rs new file mode 100644 index 0000000..9e93a11 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/ordered_partial_stream.rs
@@ -0,0 +1,352 @@ +// 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. + +//! Partial aggregate stream for ordered group input. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; +use crate::aggregates::AggregateMode; +use crate::aggregates::order::GroupOrdering; +use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; +use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; + +/// Partial aggregate stream for `InputOrderMode::Sorted` and +/// `InputOrderMode::PartiallySorted`. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// If the input is ordered by `k`, the aggregate can use ordered partial and +/// final stages: +/// +/// ## Plan +/// AggregateExec(stage=final, ordered) +/// -- RepartitionExec(hash(k), preserves_order=true) +/// ---- AggregateExec(stage=partial, ordered) +/// +/// ## Partial Stage Behavior +/// Input: raw rows +/// Output: partial states for all groups (for example, `AVG(x)` emits `SUM(x)` +/// and `COUNT(x)`) +/// +/// ## Final Stage Behavior +/// Input: partial states +/// Output: results for all groups (for example, `AVG(x)` calculated from the +/// state) +/// +/// # Order-based Optimization +/// +/// For the aggregation work, the hash aggregation implementation is reused. +/// +/// After each input batch, check whether any groups can be emitted eagerly to +/// improve memory efficiency. For example, if the last group key seen is +/// `k = 100`, it is safe to emit all groups with keys less than 100 because the +/// input is ordered. +/// +/// # Memory Pressure and Spilling +/// +/// ## Fully ordered case +/// +/// If the input is ordered by every group key, for example: +/// +/// - Input order: `a, b` +/// - `GROUP BY`: `a, b` +/// +/// Completed groups can be emitted as soon as the next group is observed. Thus, +/// only the current group remains active after completed groups are emitted, and +/// memory usage does not grow with the total number of groups. +/// +/// If a memory reservation nevertheless fails, the stream returns the error +/// directly, indicating an unexpected behavior. +/// +/// ## Partially ordered case +/// +/// If the input is ordered by only a subset of the group keys, for example: +/// +/// - Input order: `a` +/// - `GROUP BY`: `a, b` +/// +/// If one `a` value contains many distinct `b` values, the table may accumulate +/// enough groups to exceed the memory limit. +/// +/// - `OrderedPartialAggregateStream`: On reservation failure, it emits all current +/// intermediate states downstream and resets the table. The final stage can +/// merge repeated `(a, b)` state rows, so no disk spill is required. +/// - `OrderedFinalAggregateStream`: It cannot emit incomplete final results. On +/// reservation failure, it sorts the current intermediate states by the complete +/// group key and spills them as one run. After the input ends, it spills any +/// remaining states, performs a sort-preserving merge of all runs, and feeds the +/// merged input into a fully ordered final aggregate stream. +/// +/// ## Implementation Note +/// +/// This is intentionally kept simple and closely maps to +/// `GroupedHashAggregateStream` to finish the refactor sooner. +/// +/// See issue for details: <https://github.com/apache/datafusion/issues/22710> +/// +pub(crate) struct OrderedPartialAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + reduction_factor: metrics::RatioMetrics, + table: Option<OrderedAggregateTable<PartialMarker>>, +} + +impl OrderedPartialAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug_assert_eq!(agg.mode, AggregateMode::Partial); + debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reduction_factor = MetricBuilder::new(&agg.metrics) + .with_type(metrics::MetricType::Summary) + .ratio_metrics("reduction_factor", partition); + + let table = OrderedAggregateTable::<PartialMarker>::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + let reservation = + MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]")) + .with_can_spill(matches!( + table.group_ordering(), + GroupOrdering::Partial(_) + )) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + reduction_factor, + table: Some(table), + }) + } + + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema_clone = Arc::clone(&self.schema); + + let cloned_metrics = self.baseline_metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + + /// Entry point for the ordered partial aggregate state machine. + /// + /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas. + /// + /// State transitions are implemented using the generator pattern; see the comments in [`async_try_stream`]. + /// + /// Conceptual state-transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered input and aggregating batches + /// into the ordered partial aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one input batch. If the ordering proves some groups are + /// complete, yield one partial-state batch immediately, then continue + /// reading input. Otherwise continue directly with the next input batch. + /// -> DrainingFinal + /// Input was exhausted. Mark the table input as done so every remaining + /// group is safe to emit. + /// + /// DrainingFinal + /// -> DrainingFinal + /// One remaining partial-state batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> { + async_try_stream(|mut emitter| async move { + let mut table = self + .table + .take() + .expect("OrderedPartialAggregateStream state should not be None"); + + self.handle_reading_input(&mut table, &mut emitter).await?; + + // Input has exhausted, move to the final draining stage. + self.close_input(); + table.input_done(); + + self.handle_draining_final(table, &mut emitter).await?; + + Ok(()) + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + /// Consumes one ordered input batch, then immediately emits completed groups + /// if the ordering proves any group is ready. + /// + /// See comments at [`Self::create_stream`] for details. + async fn handle_reading_input( + &mut self, + table: &mut OrderedAggregateTable<PartialMarker>, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + + while let Some(batch) = self.input.next().await.transpose()? { + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + + let timer = elapsed_compute.timer(); + + table.aggregate_batch(&batch)?; + + // Check memory reservation. See function comments for details. + if let Some(batch) = self.resize_or_take_state_batch(table)? { + self.reduction_factor.add_part(batch.num_rows()); + drop(timer); + emitter.emit(batch).await; + continue; + } + + let Some(batch) = table.next_output_batch()? else { + // Can't do early emit, continue aggregating. + continue; + }; + + self.reduction_factor.add_part(batch.num_rows()); + self.reservation.try_resize(table.memory_size())?; + + drop(timer); + emitter.emit(batch).await; + } + + Ok(()) + } + + /// Update the memory reservation, and: + /// - If memory reservation succeed, returns `Ok(None)` + /// - If memory reservation failed, + /// - If input is partially ordered, materialize all the output, and + /// directly send them to the final aggregation stage. + /// Returns `Ok(Some(batch))` + /// - If input is fully ordered, directly return error. It's not + /// expected to use more than constant memory. + /// Returns `Err(..)` + /// + /// # Implementation Note + /// Incrementally output it after the blocked state management is ready, keep + /// it simple for now. + /// + /// Issue: <https://github.com/apache/datafusion/issues/7065> + fn resize_or_take_state_batch( + &mut self, + table: &mut OrderedAggregateTable<PartialMarker>, + ) -> Result<Option<RecordBatch>> { + let oom = match self.reservation.try_resize(table.memory_size()) { + Ok(()) => return Ok(None), + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Err(e), + }; + + if matches!(table.group_ordering(), GroupOrdering::Full(_)) { + return Err(oom); + } + + let Some(batch) = table.take_state_batch()? else { + return Err(oom); + }; + self.reservation.try_resize(table.memory_size())?; + Ok(Some(batch)) + } + + /// Emits one batch after input is exhausted. + /// + /// `table.input_done()` has already made every remaining group safe to emit, + /// so this state keeps draining until the table is empty. + /// + /// See comments at [`Self::create_stream`] for details. + /// + async fn handle_draining_final( + &mut self, + mut table: OrderedAggregateTable<PartialMarker>, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + while let Some(batch) = table.next_output_batch()? { + self.reduction_factor.add_part(batch.num_rows()); + + if table.is_empty() { + // Clear memory before emitting last batch so we don't have to wait for next poll to clear + drop(table); + let _ = self.reservation.try_resize(0); + drop(timer); + + emitter.emit(batch).await; + + return Ok(()); + } + + self.reservation.try_resize(table.memory_size())?; + + timer.done(); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + // was empty + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs new file mode 100644 index 0000000..2f4535e --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/partial_reduce_stream.rs
@@ -0,0 +1,385 @@ +// 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. + +//! Partial-reduce hash aggregation stream implementation. +//! +//! This stream is part of the incremental migration from +//! [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! See issue for details: <https://github.com/apache/datafusion/issues/22710> + +use std::ops::ControlFlow; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, PartialReduceMarker}; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can combine multiple partial stages before final +/// evaluation. This stream implements the partial-reduce stage. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=final) +/// -- RepartitionExec(hash(k)) +/// ---- AggregateExec(stage=partial_reduce) +/// ------ RepartitionExec(hash(k)) +/// -------- AggregateExec(stage=partial) +/// +/// Note: the example plan is only intended to demonstrate this stream's semantics; +/// the default DataFusion SQL planner does not produce plans in this shape. +/// +/// This stream implements the middle partial-reduce aggregation in the plan above. +/// +/// The motivation is to reduce shuffling traffic in a distributed setting. See +/// <https://github.com/datafusion-contrib/datafusion-distributed/issues/360> +/// +/// ## Partial-Reduce Stage Behavior +/// Input: partial aggregate state rows +/// Output: merged partial aggregate state rows +/// +/// 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. +pub(crate) struct PartialReduceHashAggregateStream { + /// Output schema: group columns followed by partial aggregate state columns. + schema: SchemaRef, + + /// Input batches containing partial aggregate state rows. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Memory reservation for group keys and accumulators. + reservation: MemoryReservation, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option<PartialReduceHashAggregateState>, +} + +/// States for partial-reduce hash aggregation processing. +// The typestate pattern mirrors the final stream and keeps the input/output +// semantics explicit for this mode. +enum PartialReduceHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable<PartialReduceMarker>, + }, + ProducingOutput { + hash_table: AggregateHashTable<PartialReduceMarker>, + }, + Done, +} + +type PartialReduceHashAggregatePoll = Poll<Option<Result<RecordBatch>>>; +type PartialReduceHashAggregateStateTransition = ControlFlow< + ( + PartialReduceHashAggregatePoll, + PartialReduceHashAggregateState, + ), + PartialReduceHashAggregateState, +>; + +impl PartialReduceHashAggregateState { + fn hash_table(&self) -> &AggregateHashTable<PartialReduceMarker> { + match self { + Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { + hash_table + } + Self::Done => unreachable!("Done state does not hold a hash table"), + } + } + + fn hash_table_mut(&mut self) -> &mut AggregateHashTable<PartialReduceMarker> { + match self { + Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { + hash_table + } + Self::Done => unreachable!("Done state does not hold a hash table"), + } + } + + fn into_hash_table(self) -> AggregateHashTable<PartialReduceMarker> { + match self { + Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { + hash_table + } + Self::Done => unreachable!("Done state does not hold a hash table"), + } + } + + fn into_producing_output(self) -> Self { + Self::ProducingOutput { + hash_table: self.into_hash_table(), + } + } + + fn into_done(self) -> Self { + Self::Done + } +} + +impl PartialReduceHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug_assert_eq!(agg.mode, super::AggregateMode::PartialReduce); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + + let hash_table = AggregateHashTable::<PartialReduceMarker>::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + + let reservation = + MemoryConsumer::new(format!("PartialReduceHashAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), + }) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable<PartialReduceMarker>, + ) -> Result<()> { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + hash_table.start_output() + } + + /// Handle ReadingInput state - aggregate partial state batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + debug_assert!(matches!( + &original_state, + PartialReduceHashAggregateState::ReadingInput { .. } + )); + debug_assert!(original_state.hash_table().is_building()); + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), + // Get a new input batch, aggregate it in the hash table + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = original_state.hash_table_mut().aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + original_state, + )); + } + + 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)) + } + // Input ends, move to output state + Poll::Ready(None) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = self.start_output(original_state.hash_table_mut()); + timer.done(); + + match result { + Ok(()) => { + ControlFlow::Continue(original_state.into_producing_output()) + } + Err(e) => { + ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + } + } + } + } + } + + /// Handle ProducingOutput state - emit merged partial aggregate state batches. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + debug_assert!(matches!( + &original_state, + PartialReduceHashAggregateState::ProducingOutput { .. } + )); + debug_assert!(!original_state.hash_table().is_building()); + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = original_state.hash_table_mut().next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let _ = self + .reservation + .try_resize(original_state.hash_table().memory_size()); + debug_assert!(batch.num_rows() > 0); + let next_state = if original_state.hash_table().is_done() { + original_state.into_done() + } else { + original_state + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Ok(None) => { + let _ = self.reservation.try_resize(0); + ControlFlow::Continue(original_state.into_done()) + } + Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + } + } +} + +impl Stream for PartialReduceHashAggregateStream { + type Item = Result<RecordBatch>; + + /// Entry point for the partial-reduce hash aggregate state machine. + /// + /// See comments in [`PartialReduceHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling partial-state input and merging those + /// states into the partial-reduce hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one partial-state input batch, update the inner aggregate + /// hash table, and continue with the next input batch. + /// + /// -> ProducingOutput + /// Input was exhausted. Move to the next state to start outputting + /// merged partial aggregate states. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One merged partial-state output batch was yielded; repeat to + /// continue producing output incrementally. + /// + /// -> Done + /// All merged partial-state output was emitted. + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + loop { + let cur_state = self + .state + .take() + .expect("PartialReduceHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ PartialReduceHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ PartialReduceHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ PartialReduceHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for PartialReduceHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs b/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs new file mode 100644 index 0000000..c6f25dc --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/single_stream.rs
@@ -0,0 +1,833 @@ +// 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. + +//! Single-stage hash aggregation stream implementation. +//! +//! This stream is part of the incremental migration from +//! [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! See issue for details: <https://github.com/apache/datafusion/issues/22710> + +use std::ops::ControlFlow; +use std::sync::Arc; +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_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::aggregate_hash_table::{AggregateHashTable, SingleMarker}; +use super::group_values::GroupByMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::{AggregateExec, create_schema}; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can run the full logical aggregation in one operator. This +/// stream implements the single stage for grouped hash aggregation. +/// +/// This aggregation variant is useful when: +/// - There is only one partition (config `target_partitions` is set to 1) +/// - When input is already partitioned (`t` is backed by Parquet files, that is range/hash +/// partitioned on the group keys), the single aggregation mode is the most efficient +/// approach to use. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=single) +/// -- DataSourceExec(t) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final aggregate values for all groups (for example, `AVG(x)`) +/// +/// This stream implements the complete aggregation without a partial/final +/// split. It consumes raw input rows and emits final aggregate values. +/// +/// # Spilling +/// +/// During aggregation, group keys and states accumulate. If memory usage exceeds +/// the budget, spilling is triggered as follows: +/// 1. After aggregating a new input batch, if the memory reservation exceeds its +/// limit, spill all accumulated groups and states. +/// - Sort all groups by the group keys before spilling. +/// 2. Repeat until the input is exhausted. +/// 3. Perform a sort-preserving merge of all spill files and feed the merged output +/// into an ordered streaming aggregation, which ensures bounded memory usage and +/// evaluates the final result. +/// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +pub(crate) struct SingleHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Memory reservation for group keys, accumulators, and spill sorting. + reservation: MemoryReservation, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option<SingleHashAggregateState>, +} + +/// Spill configuration and accumulated runs for single hash aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct SingleSpillContext { + /// Aggregate configuration used to construct the final replay stream. + /// + /// Spilled rows already contain evaluated group keys and intermediate + /// aggregate states. Replay must therefore use final aggregation semantics + /// and column-based group expressions rather than evaluating the raw input + /// expressions a second time. After the spill files are merged into ordered + /// input, this configuration is used to construct an + /// [`OrderedFinalAggregateStream`], and perform the final evaluation step. + final_agg: AggregateExec, + /// Task context. + context: Arc<TaskContext>, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Full group-key ordering kept by every spill file and the merged input. + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Spill runs waiting to be merged, they're all sorted by full group-by keys. + spills: Vec<SortedSpillFile>, +} + +/// See comments at `poll_next()` for details. +enum SingleHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable<SingleMarker>, + spill_context: Option<Box<SingleSpillContext>>, + }, + Spilling { + hash_table: AggregateHashTable<SingleMarker>, + spill_context: Box<SingleSpillContext>, + }, + ProducingOutput { + hash_table: AggregateHashTable<SingleMarker>, + }, + PreparingMergeInput { + hash_table: AggregateHashTable<SingleMarker>, + spill_context: Box<SingleSpillContext>, + }, + 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, +} + +type SingleHashAggregatePoll = Poll<Option<Result<RecordBatch>>>; +type SingleHashAggregateStateTransition = ControlFlow< + (SingleHashAggregatePoll, SingleHashAggregateState), + SingleHashAggregateState, +>; + +impl SingleSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + batch_size: usize, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result<Self> { + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + let output_ordering = agg.cache.output_ordering(); + let spill_sort_exprs = + group_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Single hash aggregate spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + // See `SingleSpillContext::final_agg` comments for `final_agg`'s usage + let mut final_agg = agg.clone(); + final_agg.mode = match agg.mode { + AggregateMode::Single => AggregateMode::Final, + AggregateMode::SinglePartitioned => AggregateMode::FinalPartitioned, + mode => { + return internal_err!( + "Single hash aggregate spill cannot replay aggregate mode {mode:?}" + ); + } + }; + final_agg.group_by = Arc::new(agg.group_by.as_final()); + final_agg.input_order_mode = InputOrderMode::Sorted; + + Ok(Self { + final_agg, + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`SingleHashAggregateStream`] for spilling details. + fn spill_table( + &mut self, + hash_table: &mut AggregateHashTable<SingleMarker>, + ) -> Result<()> { + let Some(batch) = hash_table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "SingleHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Single hash aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run, and do the aggregate evaluation with + /// [`OrderedFinalAggregateStream`] + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result<SendableRecordBatchStream> { + let Self { + final_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &final_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl SingleHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc<TaskContext>, + partition: usize, + ) -> Result<Self> { + debug_assert!(matches!( + agg.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + )); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let state_schema = Arc::new(create_schema( + input_schema.as_ref(), + &agg.group_by, + &agg.aggr_expr, + AggregateMode::Partial, + )?); + + let hash_table = AggregateHashTable::<SingleMarker>::new( + agg, + partition, + Arc::clone(&schema), + Arc::clone(&state_schema), + batch_size, + )?; + + let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + Some(Box::new(SingleSpillContext::new( + agg, + context, + partition, + batch_size, + &state_schema, + spill_metrics, + )?)) + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(SingleHashAggregateState::ReadingInput { + hash_table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> SingleHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + SingleHashAggregateState::Error, + )) + } + + fn break_with_internal_err(message: &str) -> SingleHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + hash_table: &AggregateHashTable<SingleMarker>, + spill_context: Option<&SingleSpillContext>, + ) -> usize { + let table_size = hash_table.memory_size(); + if spill_context.is_some() { + // See `SingleHashAggregateStream` comments for how this is estimated. + table_size.saturating_add( + hash_table + .building_group_count() + .saturating_mul(size_of::<u32>()), + ) + } else { + table_size + } + } + + /// Consumes one raw input batch and updates the single-stage hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::ReadingInput { + mut hash_table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected ReadingInput state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + SingleHashAggregateState::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(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + 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(_)) => { + let Some(spill_context) = spill_context else { + return Self::break_with_err(e.context( + "Single hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Single hash aggregate ran out of memory with no aggregated groups", + ); + } + return ControlFlow::Continue( + SingleHashAggregateState::Spilling { + hash_table, + spill_context, + }, + ); + } + Err(e) => { + return Self::break_with_err(e); + } + } + + ControlFlow::Continue(SingleHashAggregateState::ReadingInput { + hash_table, + spill_context, + }) + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + SingleHashAggregateState::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( + SingleHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + } + } + } + } + } + } + + /// 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( + &mut self, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::Spilling { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Single 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( + "Single hash aggregation entered Spilling with an empty table", + ); + } + + 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")); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input. + Ok(()) => ControlFlow::Continue(SingleHashAggregateState::ReadingInput { + hash_table, + spill_context: Some(spill_context), + }), + Err(e) => Self::break_with_err(e), + } + } + + /// 1. Spills the last in-memory run. + /// 2. 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 + /// 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( + &mut self, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::PreparingMergeInput { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Single 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 group_by_metrics = hash_table.group_by_metrics().clone(); + drop(hash_table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ), + Err(e) => Err(e), + } + } + Err(e) => Err(e), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(SingleHashAggregateState::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: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::MergingSpills { mut stream } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + SingleHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + SingleHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => ControlFlow::Continue(SingleHashAggregateState::Done), + } + } + + /// Emits one batch after input is exhausted. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::ProducingOutput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected ProducingOutput state", + ); + }; + + 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); + } + SingleHashAggregateState::Done + } else { + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) + { + return Self::break_with_err(e); + } + SingleHashAggregateState::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 = SingleHashAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + ControlFlow::Continue(next_state) + } + } + } +} + +impl Stream for SingleHashAggregateStream { + type Item = Result<RecordBatch>; + + /// Entry point for the single hash aggregate state machine. + /// + /// See comments in [`SingleHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling raw input rows and aggregating those + /// rows into the single-stage hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one raw 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. Start outputting final 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<Option<Self::Item>> { + loop { + let cur_state = self + .state + .take() + .expect("SingleHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ SingleHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ SingleHashAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ SingleHashAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ SingleHashAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } + state @ SingleHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ SingleHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } + state @ SingleHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, SingleHashAggregateState::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(SingleHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for SingleHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/skip_partial.rs b/native/vendor/datafusion-physical-plan/src/aggregates/skip_partial.rs new file mode 100644 index 0000000..20e17d2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/skip_partial.rs
@@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::record_batch::RecordBatch; + +use crate::metrics; + +/// Tracks if the aggregate should skip partial aggregations +/// +/// See "partial aggregation" discussion on +/// [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +pub(super) struct SkipAggregationProbe { + // ======================================================================== + // PROPERTIES: + // These fields are initialized at the start and remain constant throughout + // the execution. + // ======================================================================== + /// Aggregation ratio check performed when the number of input rows exceeds + /// this threshold (from `SessionConfig`) + probe_rows_threshold: usize, + /// Maximum ratio of `num_groups` to `input_rows` for continuing aggregation + /// (from `SessionConfig`). If the ratio exceeds this value, aggregation + /// is skipped and input rows are directly converted to output + probe_ratio_threshold: f64, + + // ======================================================================== + // STATES: + // Fields changes during execution. Can be buffer, or state flags that + // influence the execution in parent `GroupedHashAggregateStream` + // ======================================================================== + /// Number of processed input rows (updated during probing) + input_rows: usize, + /// Number of total group values for `input_rows` (updated during probing) + num_groups: usize, + + /// Flag indicating further data aggregation may be skipped (decision made + /// when probing complete) + should_skip: bool, + /// Flag indicating further updates of `SkipAggregationProbe` state won't + /// make any effect (set either while probing or on probing completion) + is_locked: bool, + + // ======================================================================== + // METRICS: + // ======================================================================== + /// Number of rows where state was output without aggregation. + /// + /// * If 0, all input rows were aggregated (should_skip was always false) + /// + /// * if greater than zero, the number of rows which were output directly + /// without aggregation + skipped_aggregation_rows: metrics::Count, +} + +impl SkipAggregationProbe { + pub(super) fn new( + probe_rows_threshold: usize, + probe_ratio_threshold: f64, + skipped_aggregation_rows: metrics::Count, + ) -> Self { + Self { + input_rows: 0, + num_groups: 0, + probe_rows_threshold, + probe_ratio_threshold, + should_skip: false, + is_locked: false, + skipped_aggregation_rows, + } + } + + /// Updates `SkipAggregationProbe` state: + /// - increments the number of input rows + /// - replaces the number of groups with the new value + /// - on `probe_rows_threshold` exceeded calculates + /// aggregation ratio and sets `should_skip` flag + /// - if `should_skip` is set, locks further state updates + pub(super) fn update_state(&mut self, input_rows: usize, num_groups: usize) { + if self.is_locked { + return; + } + self.input_rows += input_rows; + self.num_groups = num_groups; + if self.input_rows >= self.probe_rows_threshold { + self.should_skip = self.num_groups as f64 / self.input_rows as f64 + > self.probe_ratio_threshold; + // Set is_locked to true only if we have decided to skip, otherwise we can try to skip + // during processing the next record_batch. + self.is_locked = self.should_skip; + } + } + + pub(super) fn should_skip(&self) -> bool { + self.should_skip + } + + /// Record the number of rows that were output directly without aggregation + pub(super) fn record_skipped(&mut self, batch: &RecordBatch) { + self.skipped_aggregation_rows.add(batch.num_rows()); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream; + use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use crate::execution_plan::ExecutionPlan; + use crate::test::TestMemoryExec; + + use std::sync::Arc; + + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_execution::TaskContext; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use futures::StreamExt; + + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_skip_aggregation_probe_not_locked_until_skip() -> Result<()> { + // Test that the probe is not locked until we actually decide to skip. + // This allows us to continue evaluating the skip condition across multiple batches. + // + // Scenario: + // - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip + // - Batch 2: Now hits ratio threshold (high cardinality) -> skip + // + // Without the fix, the probe would be locked after batch 1, preventing the skip + // decision from being made on batch 2. + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int32, false), + ])); + + // Configure thresholds: + // - probe_rows_threshold: 100 rows + // - probe_ratio_threshold: 0.8 (80%) + let probe_rows_threshold = 100; + let probe_ratio_threshold = 0.8; + + // Batch 1: 100 rows with only 10 unique groups + // Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip + // This will hit the rows threshold but not the ratio threshold + let batch1_rows = 100; + let batch1_groups = 10; + let mut group_ids_batch1 = Vec::new(); + for i in 0..batch1_rows { + group_ids_batch1.push((i % batch1_groups) as i32); + } + let values_batch1: Vec<i32> = vec![1; batch1_rows]; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch1)), + Arc::new(Int32Array::from(values_batch1)), + ], + )?; + + // Batch 2: 360 rows with 360 unique NEW groups (starting from group 10) + // After batch 2, total: 460 rows, 370 groups + // Ratio: 370/460 is about 0.804 (80.4%) > 0.8 -> SHOULD decide to skip + let batch2_rows = 360; + let batch2_groups = 360; + let group_ids_batch2: Vec<i32> = (batch1_groups..(batch1_groups + batch2_groups)) + .map(|x| x as i32) + .collect(); + let values_batch2: Vec<i32> = vec![1; batch2_rows]; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch2)), + Arc::new(Int32Array::from(values_batch2)), + ], + )?; + + // Batch 3: This batch should be skipped since we decided to skip after batch 2 + // 100 rows with 100 unique groups (continuing from where batch 2 left off) + let batch3_rows = 100; + let batch3_groups = 100; + let batch3_start_group = batch1_groups + batch2_groups; + let group_ids_batch3: Vec<i32> = (batch3_start_group + ..(batch3_start_group + batch3_groups)) + .map(|x| x as i32) + .collect(); + let values_batch3: Vec<i32> = vec![1; batch3_rows]; + + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch3)), + Arc::new(Int32Array::from(values_batch3)), + ], + )?; + + let input_partitions = vec![vec![batch1, batch2, batch3]]; + + let runtime = RuntimeEnvBuilder::default().build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure skip aggregation settings + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + GroupedHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Check that skip aggregation actually happened. + // The key metric is skipped_aggregation_rows. + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + // We expect batch 3's rows to be skipped (100 rows) + assert_eq!( + skipped_rows, batch3_rows, + "Expected batch 3's rows ({batch3_rows}) to be skipped", + ); + + Ok(()) + } + + #[test] + fn test_skip_aggregation_probe_equality_does_not_skip() { + // When num_groups / input_rows == probe_ratio_threshold, the `>` boundary + // means we must NOT skip: equality is not sufficient to trigger skip. + let threshold_ratio = 0.5_f64; + let threshold_rows = 10_usize; + let mut probe = SkipAggregationProbe::new( + threshold_rows, + threshold_ratio, + metrics::Count::new(), + ); + + // 10 rows, 5 groups: ratio = 5/10 = 0.5 exactly equals threshold + probe.update_state(10, 5); + + assert!( + !probe.should_skip(), + "ratio == threshold should not trigger skip (boundary is exclusive)" + ); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/topk/hash_table.rs b/native/vendor/datafusion-physical-plan/src/aggregates/topk/hash_table.rs new file mode 100644 index 0000000..adc8f8c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/topk/hash_table.rs
@@ -0,0 +1,727 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A wrapper around `hashbrown::HashTable` that allows entries to be tracked by index + +use crate::aggregates::group_values::HashValue; +use crate::aggregates::topk::heap::Comparable; +use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, LargeStringArray, PrimitiveArray, StringArray, + StringViewArray, builder::PrimitiveBuilder, cast::AsArray, downcast_primitive, +}; +use arrow::datatypes::{DataType, i256}; +use datafusion_common::Result; +use datafusion_common::exec_datafusion_err; +use datafusion_common::hash_utils::RandomState; +use half::f16; +use hashbrown::hash_table::HashTable; +use std::fmt::Debug; +use std::hash::BuildHasher; +use std::sync::Arc; + +/// A "type alias" for Keys which are stored in our map +pub trait KeyType: Clone + Comparable + Debug {} + +impl<T> KeyType for T where T: Clone + Comparable + Debug {} + +/// `heap_idx` assigned to groups whose aggregate values are all NULL. Such +/// groups are tracked in the hash table only (they never enter the heap), so +/// they can be emitted with a NULL aggregate value at the end. +const NULL_HEAP_IDX: usize = usize::MAX; + +/// An entry in our hash table that: +/// 1. memoizes the hash +/// 2. contains the key (ID) +/// 3. contains the value (heap_idx - an index into the corresponding heap) +pub struct HashTableItem<ID: KeyType> { + hash: u64, + pub id: ID, + pub heap_idx: usize, +} + +/// A custom wrapper around `hashbrown::HashTable` that: +/// 1. limits the number of entries to the top K +/// 2. Allocates a capacity greater than top K to maintain a low-fill factor and prevent resizing +/// 3. Tracks indexes to allow corresponding heap to refer to entries by index vs hash +struct TopKHashTable<ID: KeyType> { + map: HashTable<usize>, + // Store the actual items separately to allow for index-based access + store: Vec<Option<HashTableItem<ID>>>, + // Free indexes in the store for reuse + free_indices: Vec<usize>, + // The maximum number of entries allowed + limit: usize, + // Number of entries registered as all-NULL (heap_idx == NULL_HEAP_IDX) + null_count: usize, +} + +/// Outcome of [`ArrowHashTable::find_or_insert`], letting the caller keep its +/// own all-NULL group accounting in sync without an extra lookup. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InsertKind { + /// The group already existed as a valued group + Existing, + /// The group was newly inserted as a valued group + New, + /// The group was registered as all-NULL and has now been converted into a + /// valued group + ReplacedNull, +} + +/// An interface to hide the generic type signature of TopKHashTable behind arrow arrays +pub trait ArrowHashTable { + fn set_batch(&mut self, ids: ArrayRef); + fn len(&self) -> usize; + fn update_heap_idx(&mut self, mapper: &[(usize, usize)]); + fn heap_idx_at(&self, map_idx: usize) -> usize; + fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef; + fn find_or_insert( + &mut self, + row_idx: usize, + replace_idx: usize, + ) -> (usize, InsertKind); + /// Register the group at `row_idx` as all-NULL. Returns true if it was + /// newly registered; false if the group is already tracked or the NULL + /// group limit has been reached. + fn insert_null(&mut self, row_idx: usize) -> bool; + /// Remove the group at `row_idx` if it is registered as all-NULL. Returns + /// true if a NULL registration was removed. + fn remove_if_null(&mut self, row_idx: usize) -> bool; + /// Store indexes of all groups registered as all-NULL + fn null_map_idxs(&self) -> Vec<usize>; +} + +/// Returns true if the given data type can be used as a top-K aggregation hash key. +/// +/// Supported types include Arrow primitives (integers, floats, decimals, intervals) +/// and UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`). This is used internally by +/// `PriorityMap::supports()` to validate grouping key type compatibility. +pub fn is_supported_hash_key_type(kt: &DataType) -> bool { + kt.is_primitive() + || matches!( + kt, + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 + ) +} + +// An implementation of ArrowHashTable for String keys +pub struct StringHashTable { + owned: ArrayRef, + map: TopKHashTable<Option<String>>, + rnd: RandomState, + data_type: DataType, +} + +// An implementation of ArrowHashTable for any `ArrowPrimitiveType` key +struct PrimitiveHashTable<VAL: ArrowPrimitiveType> +where + Option<<VAL as ArrowPrimitiveType>::Native>: Comparable, +{ + owned: ArrayRef, + map: TopKHashTable<Option<VAL::Native>>, + rnd: RandomState, + kt: DataType, +} + +impl StringHashTable { + pub fn new(limit: usize, data_type: DataType) -> Self { + let vals: Vec<&str> = Vec::new(); + let owned: ArrayRef = match data_type { + DataType::Utf8 => Arc::new(StringArray::from(vals)), + DataType::Utf8View => Arc::new(StringViewArray::from(vals)), + DataType::LargeUtf8 => Arc::new(LargeStringArray::from(vals)), + _ => panic!("Unsupported data type"), + }; + + Self { + owned, + map: TopKHashTable::new(limit, limit * 10), + rnd: RandomState::default(), + data_type, + } + } + + /// Extracts the string value at the given row index, handling nulls and different string types. + /// + /// Returns `None` if the value is null, otherwise `Some(value.to_string())`. + fn extract_string_value(&self, row_idx: usize) -> Option<String> { + let is_null_and_value = match self.data_type { + DataType::Utf8 => { + let arr = self.owned.as_string::<i32>(); + (arr.is_null(row_idx), arr.value(row_idx)) + } + DataType::LargeUtf8 => { + let arr = self.owned.as_string::<i64>(); + (arr.is_null(row_idx), arr.value(row_idx)) + } + DataType::Utf8View => { + let arr = self.owned.as_string_view(); + (arr.is_null(row_idx), arr.value(row_idx)) + } + _ => panic!("Unsupported data type"), + }; + + let (is_null, value) = is_null_and_value; + if is_null { + None + } else { + Some(value.to_string()) + } + } + + /// Computes the id and its hash for the given row, for hash table lookups + fn id_and_hash(&self, row_idx: usize) -> (Option<String>, u64) { + let id = self.extract_string_value(row_idx); + let hash = self.rnd.hash_one(id.as_deref()); + (id, hash) + } +} + +impl ArrowHashTable for StringHashTable { + fn set_batch(&mut self, ids: ArrayRef) { + self.owned = ids; + } + + fn len(&self) -> usize { + self.map.len() + } + + fn update_heap_idx(&mut self, mapper: &[(usize, usize)]) { + self.map.update_heap_idx(mapper); + } + + fn heap_idx_at(&self, map_idx: usize) -> usize { + self.map.heap_idx_at(map_idx) + } + + fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef { + let ids = self.map.take_all(indexes); + match self.data_type { + DataType::Utf8 => Arc::new(StringArray::from(ids)), + DataType::LargeUtf8 => Arc::new(LargeStringArray::from(ids)), + DataType::Utf8View => Arc::new(StringViewArray::from(ids)), + _ => unreachable!(), + } + } + + fn find_or_insert( + &mut self, + row_idx: usize, + replace_idx: usize, + ) -> (usize, InsertKind) { + let id = self.extract_string_value(row_idx); + + // Compute hash and create equality closure for hash table lookup. + let hash = self.rnd.hash_one(id.as_deref()); + let id_for_eq = id.clone(); + let eq = move |mi: &Option<String>| id_for_eq.as_deref() == mi.as_deref(); + + // Use entry API to avoid double lookup + self.map.find_or_insert(hash, id, replace_idx, eq) + } + + fn insert_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let id_for_eq = id.clone(); + let eq = move |mi: &Option<String>| id_for_eq.as_deref() == mi.as_deref(); + self.map.insert_null(hash, id, eq) + } + + fn remove_if_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let eq = move |mi: &Option<String>| id.as_deref() == mi.as_deref(); + self.map.remove_if_null(hash, eq) + } + + fn null_map_idxs(&self) -> Vec<usize> { + self.map.null_map_idxs() + } +} + +impl<VAL: ArrowPrimitiveType> PrimitiveHashTable<VAL> +where + Option<<VAL as ArrowPrimitiveType>::Native>: Comparable, + Option<<VAL as ArrowPrimitiveType>::Native>: HashValue, +{ + pub fn new(limit: usize, kt: DataType) -> Self { + let owned = Arc::new( + PrimitiveArray::<VAL>::builder(0) + .with_data_type(kt.clone()) + .finish(), + ); + Self { + owned, + map: TopKHashTable::new(limit, limit * 10), + rnd: RandomState::default(), + kt, + } + } + + /// Computes the id and its hash for the given row, for hash table lookups + fn id_and_hash(&self, row_idx: usize) -> (Option<VAL::Native>, u64) { + let ids = self.owned.as_primitive::<VAL>(); + let id: Option<VAL::Native> = if ids.is_null(row_idx) { + None + } else { + Some(ids.value(row_idx)) + }; + let hash: u64 = id.hash(&self.rnd); + (id, hash) + } +} + +impl<VAL: ArrowPrimitiveType> ArrowHashTable for PrimitiveHashTable<VAL> +where + Option<<VAL as ArrowPrimitiveType>::Native>: Comparable, + Option<<VAL as ArrowPrimitiveType>::Native>: HashValue, +{ + fn set_batch(&mut self, ids: ArrayRef) { + self.owned = ids; + } + + fn len(&self) -> usize { + self.map.len() + } + + fn update_heap_idx(&mut self, mapper: &[(usize, usize)]) { + self.map.update_heap_idx(mapper); + } + + fn heap_idx_at(&self, map_idx: usize) -> usize { + self.map.heap_idx_at(map_idx) + } + + fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef { + let ids = self.map.take_all(indexes); + let mut builder: PrimitiveBuilder<VAL> = + PrimitiveArray::builder(ids.len()).with_data_type(self.kt.clone()); + for id in ids.into_iter() { + match id { + None => builder.append_null(), + Some(id) => builder.append_value(id), + } + } + let ids = builder.finish(); + Arc::new(ids) + } + + fn find_or_insert( + &mut self, + row_idx: usize, + replace_idx: usize, + ) -> (usize, InsertKind) { + let ids = self.owned.as_primitive::<VAL>(); + let id: Option<VAL::Native> = if ids.is_null(row_idx) { + None + } else { + Some(ids.value(row_idx)) + }; + // Compute hash and create equality closure for hash table lookup. + let hash: u64 = id.hash(&self.rnd); + let eq = |mi: &Option<VAL::Native>| id == *mi; + + // Use entry API to avoid double lookup + self.map.find_or_insert(hash, id, replace_idx, eq) + } + + fn insert_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let eq = move |mi: &Option<VAL::Native>| id == *mi; + self.map.insert_null(hash, id, eq) + } + + fn remove_if_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let eq = move |mi: &Option<VAL::Native>| id == *mi; + self.map.remove_if_null(hash, eq) + } + + fn null_map_idxs(&self) -> Vec<usize> { + self.map.null_map_idxs() + } +} + +use hashbrown::hash_table::Entry; +impl<ID: KeyType + PartialEq> TopKHashTable<ID> { + pub fn new(limit: usize, capacity: usize) -> Self { + Self { + map: HashTable::with_capacity(capacity), + store: Vec::with_capacity(capacity), + free_indices: Vec::new(), + limit, + null_count: 0, + } + } + + pub fn heap_idx_at(&self, map_idx: usize) -> usize { + self.store[map_idx].as_ref().unwrap().heap_idx + } + + /// Remove the entry stored at `map_idx`, freeing its store slot for reuse + fn remove_at(&mut self, map_idx: usize) { + let item_to_remove = self.store[map_idx].as_ref().unwrap(); + let hash = item_to_remove.hash; + let id_to_remove = &item_to_remove.id; + + let eq = |&idx: &usize| self.store[idx].as_ref().unwrap().id == *id_to_remove; + let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + match self.map.entry(hash, eq, hasher) { + Entry::Occupied(entry) => { + let (removed_idx, _) = entry.remove(); + self.store[removed_idx] = None; + self.free_indices.push(removed_idx); + } + Entry::Vacant(_) => unreachable!(), + } + } + + pub fn remove_if_full(&mut self, replace_idx: usize) -> usize { + // All-NULL groups are tracked outside the heap, so only valued + // groups count towards the limit here + let valued_len = self.map.len() - self.null_count; + if valued_len >= self.limit { + self.remove_at(replace_idx); + 0 // if full, always replace top node + } else { + valued_len // if we're not full, always append to end + } + } + + fn update_heap_idx(&mut self, mapper: &[(usize, usize)]) { + for (m, h) in mapper { + self.store[*m].as_mut().unwrap().heap_idx = *h; + } + } + + /// Find an existing entry or insert a new one, avoiding double hash table lookup. + /// Returns (map_idx, kind) where kind describes whether the group already + /// existed, was newly inserted, or was converted from an all-NULL group. + /// If inserting a new entry and the table is full, replaces the entry at replace_idx. + pub fn find_or_insert( + &mut self, + hash: u64, + id: ID, + replace_idx: usize, + mut eq: impl FnMut(&ID) -> bool, + ) -> (usize, InsertKind) { + // Check if entry exists - this is the only hash table lookup + let mut replaced_null = false; + { + let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + if let Some(&map_idx) = self.map.find(hash, eq_fn) { + if self.store[map_idx].as_ref().unwrap().heap_idx == NULL_HEAP_IDX { + // This group was registered as all-NULL but now produced a + // value: unregister it so it is inserted as a valued group + self.remove_at(map_idx); + self.null_count -= 1; + replaced_null = true; + } else { + return (map_idx, InsertKind::Existing); + } + } + } + + // Entry doesn't exist - compute heap_idx and prepare item + let heap_idx = self.remove_if_full(replace_idx); + let mi = HashTableItem::new(hash, id, heap_idx); + let store_idx = if let Some(idx) = self.free_indices.pop() { + self.store[idx] = Some(mi); + idx + } else { + self.store.push(Some(mi)); + self.store.len() - 1 + }; + + // Reserve space if needed + let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + if self.map.len() == self.map.capacity() { + self.map.reserve(self.limit, hasher); + } + + // Insert without checking again since we already confirmed it doesn't exist + self.map.insert_unique(hash, store_idx, hasher); + let kind = if replaced_null { + InsertKind::ReplacedNull + } else { + InsertKind::New + }; + (store_idx, kind) + } + + /// Register a group whose aggregate values are all NULL, unless it is + /// already tracked. NULL groups are stored with a sentinel `heap_idx` and + /// never enter the heap. At most `limit` NULL groups are tracked: they all + /// tie on the sort key, so any `limit` of them is a valid top-k superset. + /// Returns true if the group was newly registered. + pub fn insert_null( + &mut self, + hash: u64, + id: ID, + mut eq: impl FnMut(&ID) -> bool, + ) -> bool { + { + let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + if self.map.find(hash, eq_fn).is_some() { + return false; + } + } + if self.null_count >= self.limit { + return false; + } + + let mi = HashTableItem::new(hash, id, NULL_HEAP_IDX); + let store_idx = if let Some(idx) = self.free_indices.pop() { + self.store[idx] = Some(mi); + idx + } else { + self.store.push(Some(mi)); + self.store.len() - 1 + }; + + let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + if self.map.len() == self.map.capacity() { + self.map.reserve(self.limit, hasher); + } + self.map.insert_unique(hash, store_idx, hasher); + self.null_count += 1; + true + } + + /// Remove the given group if it is registered as all-NULL. Used when an + /// all-NULL group produces a value that loses to the current top-k: the + /// group can no longer reach the top-k, but it must not be emitted with a + /// NULL value either. Returns true if a NULL registration was removed. + pub fn remove_if_null(&mut self, hash: u64, mut eq: impl FnMut(&ID) -> bool) -> bool { + let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + if let Some(&map_idx) = self.map.find(hash, eq_fn) + && self.store[map_idx].as_ref().unwrap().heap_idx == NULL_HEAP_IDX + { + self.remove_at(map_idx); + self.null_count -= 1; + return true; + } + false + } + + /// Store indexes of all groups registered as all-NULL + pub fn null_map_idxs(&self) -> Vec<usize> { + self.store + .iter() + .enumerate() + .filter_map(|(idx, item)| { + item.as_ref() + .filter(|item| item.heap_idx == NULL_HEAP_IDX) + .map(|_| idx) + }) + .collect() + } + + pub fn len(&self) -> usize { + self.map.len() + } + + pub fn take_all(&mut self, idxs: Vec<usize>) -> Vec<ID> { + let ids = idxs + .into_iter() + .map(|idx| self.store[idx].take().unwrap().id) + .collect(); + self.map.clear(); + self.store.clear(); + self.free_indices.clear(); + self.null_count = 0; + ids + } +} + +impl<ID: KeyType> HashTableItem<ID> { + pub fn new(hash: u64, id: ID, heap_idx: usize) -> Self { + Self { hash, id, heap_idx } + } +} + +impl HashValue for Option<String> { + fn hash(&self, state: &RandomState) -> u64 { + state.hash_one(self) + } +} + +macro_rules! hash_float { + ($($t:ty),+) => { + $(impl HashValue for Option<$t> { + fn hash(&self, state: &RandomState) -> u64 { + self.map(|me| me.hash(state)).unwrap_or(0) + } + })+ + }; +} + +macro_rules! has_integer { + ($($t:ty),+) => { + $(impl HashValue for Option<$t> { + fn hash(&self, state: &RandomState) -> u64 { + self.map(|me| me.hash(state)).unwrap_or(0) + } + })+ + }; +} + +has_integer!(i8, i16, i32, i64, i128, i256); +has_integer!(u8, u16, u32, u64); +has_integer!(IntervalDayTime, IntervalMonthDayNano); +hash_float!(f16, f32, f64); + +pub fn new_hash_table( + limit: usize, + kt: DataType, +) -> Result<Box<dyn ArrowHashTable + Send>> { + macro_rules! downcast_helper { + ($kt:ty, $d:ident) => { + return Ok(Box::new(PrimitiveHashTable::<$kt>::new(limit, kt))) + }; + } + + downcast_primitive! { + kt => (downcast_helper, kt), + DataType::Utf8 => return Ok(Box::new(StringHashTable::new(limit, DataType::Utf8))), + DataType::LargeUtf8 => return Ok(Box::new(StringHashTable::new(limit, DataType::LargeUtf8))), + DataType::Utf8View => return Ok(Box::new(StringHashTable::new(limit, DataType::Utf8View))), + _ => {} + } + + Err(exec_datafusion_err!( + "Can't create HashTable for type: {kt:?}" + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::TimestampMillisecondArray; + use arrow_schema::TimeUnit; + use std::collections::BTreeMap; + + #[test] + fn should_emit_correct_type() -> Result<()> { + let ids = + TimestampMillisecondArray::from(vec![1000]).with_timezone("UTC".to_string()); + let dt = DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())); + let mut ht = new_hash_table(1, dt.clone())?; + ht.set_batch(Arc::new(ids)); + ht.find_or_insert(0, 0); + let ids = ht.take_all(vec![0]); + assert_eq!(ids.data_type(), &dt); + + Ok(()) + } + + #[test] + fn should_resize_properly() -> Result<()> { + let mut heap_to_map = BTreeMap::<usize, usize>::new(); + // Create TopKHashTable with limit=5 and capacity=3 to force resizing + let mut map = TopKHashTable::<Option<String>>::new(5, 3); + + // Insert 5 entries, tracking the heap-to-map index mapping + for (heap_idx, id) in ["1", "2", "3", "4", "5"].iter().enumerate() { + let value = Some(id.to_string()); + let hash = heap_idx as u64; + let (map_idx, kind) = + map.find_or_insert(hash, value.clone(), heap_idx, |v| *v == value); + assert_eq!(kind, InsertKind::New, "Entry should be new"); + heap_to_map.insert(heap_idx, map_idx); + } + + // Verify all 5 entries are present + assert_eq!(map.len(), 5); + + // Verify that the hash table resized properly (capacity should have grown beyond 3) + // This is implicit - if it didn't resize, insertions would have failed or been slow + + // Drain all values in heap order + let (_heap_idxs, map_idxs): (Vec<_>, Vec<_>) = heap_to_map.into_iter().unzip(); + let ids = map.take_all(map_idxs); + + assert_eq!( + format!("{ids:?}"), + r#"[Some("1"), Some("2"), Some("3"), Some("4"), Some("5")]"# + ); + assert_eq!(map.len(), 0, "Map should have been cleared!"); + + Ok(()) + } + + #[test] + fn should_track_null_groups() -> Result<()> { + let mut map = TopKHashTable::<Option<String>>::new(2, 10); + + let a = Some("a".to_string()); + let b = Some("b".to_string()); + let c = Some("c".to_string()); + + // register two all-NULL groups; the third exceeds the NULL group limit + assert!(map.insert_null(100, a.clone(), |v| *v == a)); + assert!(map.insert_null(200, b.clone(), |v| *v == b)); + assert!(!map.insert_null(300, c.clone(), |v| *v == c)); + // re-registering an existing NULL group is a no-op + assert!(!map.insert_null(100, a.clone(), |v| *v == a)); + assert_eq!(map.null_count, 2); + assert_eq!(map.null_map_idxs(), vec![0, 1]); + + // a valued insert for a NULL group converts it to a valued group + let (map_idx, kind) = map.find_or_insert(200, b.clone(), 0, |v| *v == b); + assert_eq!(kind, InsertKind::ReplacedNull, "NULL group should convert"); + assert_eq!(map.heap_idx_at(map_idx), 0, "Heap should append at 0"); + assert_eq!(map.null_count, 1); + assert_eq!(map.null_map_idxs(), vec![0]); + + // remove the remaining NULL group; removing twice is a no-op + map.remove_if_null(100, |v| *v == a); + assert_eq!(map.null_count, 0); + assert!(map.null_map_idxs().is_empty()); + map.remove_if_null(100, |v| *v == a); + // removing a valued group via remove_if_null is a no-op + map.remove_if_null(200, |v| *v == b); + assert_eq!(map.len(), 1); + + Ok(()) + } + + #[test] + fn should_reuse_all_freed_store_slots() -> Result<()> { + let mut map = TopKHashTable::<Option<String>>::new(1, 10); + + let a = Some("a".to_string()); + let b = Some("b".to_string()); + let c = Some("c".to_string()); + + let (b_idx, kind) = map.find_or_insert(100, b.clone(), 0, |v| *v == b); + assert_eq!(kind, InsertKind::New); + assert!(map.insert_null(200, a.clone(), |v| *v == a)); + + // Converting a NULL group while the valued heap is full frees two + // slots: the NULL registration and the evicted valued group. + let (_, kind) = map.find_or_insert(200, a.clone(), b_idx, |v| *v == a); + assert_eq!(kind, InsertKind::ReplacedNull); + + // Both freed slots must remain reusable. Otherwise repeated + // conversions make the backing store grow without bound. + assert!(map.insert_null(300, c.clone(), |v| *v == c)); + assert_eq!(map.store.len(), 2); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/topk/heap.rs b/native/vendor/datafusion-physical-plan/src/aggregates/topk/heap.rs new file mode 100644 index 0000000..ca321cd --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/topk/heap.rs
@@ -0,0 +1,783 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A custom binary heap implementation for performant top K aggregation. +//! +//! the `new_heap` //! factory function selects an appropriate heap implementation +//! based on the Arrow data type. +//! +//! Supported value types include Arrow primitives (integers, floats, decimals, intervals) +//! and UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`) using lexicographic ordering. + +use arrow::array::{ArrayRef, ArrowPrimitiveType, PrimitiveArray, downcast_primitive}; +use arrow::array::{LargeStringBuilder, StringBuilder, StringViewBuilder}; +use arrow::array::{ + StringArray, + cast::AsArray, + types::{IntervalDayTime, IntervalMonthDayNano}, +}; +use arrow::buffer::ScalarBuffer; +use arrow::datatypes::{DataType, i256}; +use datafusion_common::Result; +use datafusion_common::exec_datafusion_err; + +use half::f16; +use std::cmp::Ordering; +use std::fmt::{Debug, Display, Formatter}; +use std::sync::Arc; + +/// A custom version of `Ord` that only exists to we can implement it for the Values in our heap +pub trait Comparable { + fn comp(&self, other: &Self) -> Ordering; +} + +impl Comparable for Option<String> { + fn comp(&self, other: &Self) -> Ordering { + self.cmp(other) + } +} + +/// A "type alias" for Values which are stored in our heap +pub trait ValueType: Comparable + Clone + Debug {} + +impl<T> ValueType for T where T: Comparable + Clone + Debug {} + +/// An entry in our heap, which contains both the value and a index into an external HashTable +struct HeapItem<VAL: ValueType> { + val: VAL, + map_idx: usize, +} + +/// A custom heap implementation that allows several things that couldn't be achieved with +/// `collections::BinaryHeap`: +/// 1. It allows values to be updated at arbitrary positions (when group values change) +/// 2. It can be either a min or max heap +/// 3. It can use our `HeapItem` type & `Comparable` trait +/// 4. It is specialized to grow to a certain limit, then always replace without grow & shrink +struct TopKHeap<VAL: ValueType> { + desc: bool, + len: usize, + capacity: usize, + heap: Vec<Option<HeapItem<VAL>>>, +} + +/// An interface to hide the generic type signature of TopKHeap behind arrow arrays +pub trait ArrowHeap { + fn set_batch(&mut self, vals: ArrayRef); + fn is_worse(&self, idx: usize) -> bool; + fn worst_map_idx(&self) -> usize; + fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>); + fn replace_if_better( + &mut self, + heap_idx: usize, + row_idx: usize, + map: &mut Vec<(usize, usize)>, + ); + fn drain(&mut self) -> (ArrayRef, Vec<usize>); +} + +/// An implementation of `ArrowHeap` that deals with primitive values +pub struct PrimitiveHeap<VAL: ArrowPrimitiveType> +where + <VAL as ArrowPrimitiveType>::Native: Comparable, +{ + batch: ArrayRef, + heap: TopKHeap<VAL::Native>, + desc: bool, + data_type: DataType, +} + +impl<VAL: ArrowPrimitiveType> PrimitiveHeap<VAL> +where + <VAL as ArrowPrimitiveType>::Native: Comparable, +{ + pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { + let owned: ArrayRef = Arc::new(PrimitiveArray::<VAL>::builder(0).finish()); + Self { + batch: owned, + heap: TopKHeap::new(limit, desc), + desc, + data_type, + } + } +} + +impl<VAL: ArrowPrimitiveType> ArrowHeap for PrimitiveHeap<VAL> +where + <VAL as ArrowPrimitiveType>::Native: Comparable, +{ + fn set_batch(&mut self, vals: ArrayRef) { + self.batch = vals; + } + + fn is_worse(&self, row_idx: usize) -> bool { + if !self.heap.is_full() { + return false; + } + let vals = self.batch.as_primitive::<VAL>(); + let new_val = vals.value(row_idx); + let worst_val = self.heap.worst_val().expect("Missing root"); + (!self.desc && new_val > *worst_val) || (self.desc && new_val < *worst_val) + } + + fn worst_map_idx(&self) -> usize { + self.heap.worst_map_idx() + } + + fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { + let vals = self.batch.as_primitive::<VAL>(); + let new_val = vals.value(row_idx); + self.heap.append_or_replace(new_val, map_idx, map); + } + + fn replace_if_better( + &mut self, + heap_idx: usize, + row_idx: usize, + map: &mut Vec<(usize, usize)>, + ) { + let vals = self.batch.as_primitive::<VAL>(); + let new_val = vals.value(row_idx); + self.heap.replace_if_better(heap_idx, new_val, map); + } + + fn drain(&mut self) -> (ArrayRef, Vec<usize>) { + let nulls = None; + let (vals, map_idxs) = self.heap.drain(); + let arr = PrimitiveArray::<VAL>::new(ScalarBuffer::from(vals), nulls) + .with_data_type(self.data_type.clone()); + (Arc::new(arr), map_idxs) + } +} + +/// An implementation of `ArrowHeap` that deals with string values. +/// +/// Supports all three UTF-8 string types: `Utf8`, `LargeUtf8`, and `Utf8View`. +/// String values are compared lexicographically using the compare-first pattern: +/// borrowed strings are compared before allocation, and only allocated when the +/// heap confirms they improve the top-K set. +/// +pub struct StringHeap { + batch: ArrayRef, + heap: TopKHeap<Option<String>>, + desc: bool, + data_type: DataType, +} + +impl StringHeap { + pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { + let batch: ArrayRef = Arc::new(StringArray::from(Vec::<&str>::new())); + Self { + batch, + heap: TopKHeap::new(limit, desc), + desc, + data_type, + } + } + + /// Extracts a string value from the current batch at the given row index. + /// + /// Panics if the row index is out of bounds or if the data type is not one of + /// the supported UTF-8 string types. + /// + /// Note: Null values should not appear in the input; the aggregation layer + /// ensures nulls are filtered before reaching this code. + fn value(&self, row_idx: usize) -> &str { + extract_string_value(&self.batch, &self.data_type, row_idx) + } +} + +/// Helper to extract a string value from an ArrayRef at a given index. +/// +/// Supports `Utf8`, `LargeUtf8`, and `Utf8View` data types. +/// +/// # Panics +/// Panics if the index is out of bounds or if the data type is unsupported. +fn extract_string_value<'a>( + batch: &'a ArrayRef, + data_type: &DataType, + idx: usize, +) -> &'a str { + match data_type { + DataType::Utf8 => batch.as_string::<i32>().value(idx), + DataType::LargeUtf8 => batch.as_string::<i64>().value(idx), + DataType::Utf8View => batch.as_string_view().value(idx), + _ => unreachable!("Unsupported string type: {data_type}"), + } +} + +impl ArrowHeap for StringHeap { + fn set_batch(&mut self, vals: ArrayRef) { + self.batch = vals; + } + + fn is_worse(&self, row_idx: usize) -> bool { + if !self.heap.is_full() { + return false; + } + // Compare borrowed `&str` against the worst heap value first to avoid + // allocating a `String` unless this row would actually replace an + // existing heap entry. + let new_val = self.value(row_idx); + let worst_val = self.heap.worst_val().expect("Missing root"); + match worst_val { + None => false, + Some(worst_str) => { + (!self.desc && new_val > worst_str.as_str()) + || (self.desc && new_val < worst_str.as_str()) + } + } + } + + fn worst_map_idx(&self) -> usize { + self.heap.worst_map_idx() + } + + fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { + // When appending (heap not full) we must allocate to own the string + // because it will be stored in the heap. For replacements we avoid + // allocation until `replace_if_better` confirms a replacement is + // necessary. + let new_str = self.value(row_idx).to_string(); + let new_val = Some(new_str); + self.heap.append_or_replace(new_val, map_idx, map); + } + + fn replace_if_better( + &mut self, + heap_idx: usize, + row_idx: usize, + map: &mut Vec<(usize, usize)>, + ) { + let new_str = self.value(row_idx); + let existing = self.heap.heap[heap_idx] + .as_ref() + .expect("Missing heap item"); + + // Compare borrowed reference first—no allocation yet. + // We compare the borrowed `&str` with the stored `Option<String>` and + // only allocate (`to_string()`) when a replacement is required. + match &existing.val { + None => { + // Existing is null; new value always wins + let new_val = Some(new_str.to_string()); + self.heap.replace_if_better(heap_idx, new_val, map); + } + Some(existing_str) => { + // Compare borrowed strings first + if (!self.desc && new_str < existing_str.as_str()) + || (self.desc && new_str > existing_str.as_str()) + { + let new_val = Some(new_str.to_string()); + self.heap.replace_if_better(heap_idx, new_val, map); + } + // Else: no improvement, no allocation + } + } + } + + fn drain(&mut self) -> (ArrayRef, Vec<usize>) { + let (vals, map_idxs) = self.heap.drain(); + // Use Arrow builders to safely construct arrays from the owned + // `Option<String>` values. Builders avoid needing to maintain + // references to temporary storage. + + // Macro to eliminate duplication across string builder types. + // All three builders share the same interface for append_value, + // append_null, and finish, differing only in their concrete types. + macro_rules! build_string_array { + ($builder_type:ty) => {{ + let mut builder = <$builder_type>::new(); + for val in vals { + match val { + Some(s) => builder.append_value(&s), + None => builder.append_null(), + } + } + Arc::new(builder.finish()) + }}; + } + + let arr: ArrayRef = match self.data_type { + DataType::Utf8 => build_string_array!(StringBuilder), + DataType::LargeUtf8 => build_string_array!(LargeStringBuilder), + DataType::Utf8View => build_string_array!(StringViewBuilder), + _ => unreachable!("Unsupported string type: {}", self.data_type), + }; + (arr, map_idxs) + } +} + +impl<VAL: ValueType> TopKHeap<VAL> { + pub fn new(limit: usize, desc: bool) -> Self { + Self { + desc, + capacity: limit, + len: 0, + heap: (0..=limit).map(|_| None).collect::<Vec<_>>(), + } + } + + pub fn worst_val(&self) -> Option<&VAL> { + let root = self.heap.first()?; + let hi = root.as_ref()?; + Some(&hi.val) + } + + pub fn worst_map_idx(&self) -> usize { + self.heap[0].as_ref().map(|hi| hi.map_idx).unwrap_or(0) + } + + pub fn is_full(&self) -> bool { + self.len >= self.capacity + } + + pub fn len(&self) -> usize { + self.len + } + + pub fn append_or_replace( + &mut self, + new_val: VAL, + map_idx: usize, + map: &mut Vec<(usize, usize)>, + ) { + if self.is_full() { + self.replace_root(new_val, map_idx, map); + } else { + self.append(new_val, map_idx, map); + } + } + + fn append(&mut self, new_val: VAL, map_idx: usize, mapper: &mut Vec<(usize, usize)>) { + let hi = HeapItem::new(new_val, map_idx); + self.heap[self.len] = Some(hi); + self.heapify_up(self.len, mapper); + self.len += 1; + } + + fn pop(&mut self, map: &mut Vec<(usize, usize)>) -> Option<HeapItem<VAL>> { + if self.len() == 0 { + return None; + } + if self.len() == 1 { + self.len = 0; + return self.heap[0].take(); + } + self.swap(0, self.len - 1, map); + let former_root = self.heap[self.len - 1].take(); + self.len -= 1; + self.heapify_down(0, map); + former_root + } + + pub fn drain(&mut self) -> (Vec<VAL>, Vec<usize>) { + let mut map = Vec::with_capacity(self.len); + let mut vals = Vec::with_capacity(self.len); + let mut map_idxs = Vec::with_capacity(self.len); + while let Some(worst_hi) = self.pop(&mut map) { + vals.push(worst_hi.val); + map_idxs.push(worst_hi.map_idx); + } + vals.reverse(); + map_idxs.reverse(); + (vals, map_idxs) + } + + fn replace_root( + &mut self, + new_val: VAL, + map_idx: usize, + mapper: &mut Vec<(usize, usize)>, + ) { + let hi = self.heap[0].as_mut().expect("No root"); + hi.val = new_val; + hi.map_idx = map_idx; + self.heapify_down(0, mapper); + } + + pub fn replace_if_better( + &mut self, + heap_idx: usize, + new_val: VAL, + mapper: &mut Vec<(usize, usize)>, + ) { + let existing = self.heap[heap_idx].as_mut().expect("Missing heap item"); + if (!self.desc && new_val.comp(&existing.val) != Ordering::Less) + || (self.desc && new_val.comp(&existing.val) != Ordering::Greater) + { + return; + } + existing.val = new_val; + self.heapify_down(heap_idx, mapper); + } + + fn heapify_up(&mut self, mut idx: usize, mapper: &mut Vec<(usize, usize)>) { + let desc = self.desc; + while idx != 0 { + let parent_idx = (idx - 1) / 2; + let node = self.heap[idx].as_ref().expect("No heap item"); + let parent = self.heap[parent_idx].as_ref().expect("No heap item"); + if (!desc && node.val.comp(&parent.val) != Ordering::Greater) + || (desc && node.val.comp(&parent.val) != Ordering::Less) + { + return; + } + self.swap(idx, parent_idx, mapper); + idx = parent_idx; + } + } + + fn swap(&mut self, a_idx: usize, b_idx: usize, mapper: &mut Vec<(usize, usize)>) { + let a_hi = self.heap[a_idx].take().expect("Missing heap entry"); + let b_hi = self.heap[b_idx].take().expect("Missing heap entry"); + + mapper.push((a_hi.map_idx, b_idx)); + mapper.push((b_hi.map_idx, a_idx)); + + self.heap[a_idx] = Some(b_hi); + self.heap[b_idx] = Some(a_hi); + } + + fn heapify_down(&mut self, node_idx: usize, mapper: &mut Vec<(usize, usize)>) { + let left_child = node_idx * 2 + 1; + let desc = self.desc; + let entry = self.heap.get(node_idx).expect("Missing node!"); + let entry = entry.as_ref().expect("Missing node!"); + let mut best_idx = node_idx; + let mut best_val = &entry.val; + for child_idx in left_child..=left_child + 1 { + if let Some(Some(child)) = self.heap.get(child_idx) + && ((!desc && child.val.comp(best_val) == Ordering::Greater) + || (desc && child.val.comp(best_val) == Ordering::Less)) + { + best_val = &child.val; + best_idx = child_idx; + } + } + if best_val.comp(&entry.val) != Ordering::Equal { + self.swap(best_idx, node_idx, mapper); + self.heapify_down(best_idx, mapper); + } + } + + fn _tree_print(&self, idx: usize, prefix: &str, is_tail: bool, output: &mut String) { + if let Some(Some(hi)) = self.heap.get(idx) { + let connector = if idx != 0 { + if is_tail { "└── " } else { "├── " } + } else { + "" + }; + output.push_str(&format!( + "{}{}val={:?} idx={}, bucket={}\n", + prefix, connector, hi.val, idx, hi.map_idx + )); + let new_prefix = if is_tail { "" } else { "│ " }; + let child_prefix = format!("{prefix}{new_prefix}"); + + let left_idx = idx * 2 + 1; + let right_idx = idx * 2 + 2; + + let left_exists = left_idx < self.len; + let right_exists = right_idx < self.len; + + if left_exists { + self._tree_print(left_idx, &child_prefix, !right_exists, output); + } + if right_exists { + self._tree_print(right_idx, &child_prefix, true, output); + } + } + } +} + +impl<VAL: ValueType> Display for TopKHeap<VAL> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let mut output = String::new(); + if !self.heap.is_empty() { + self._tree_print(0, "", true, &mut output); + } + write!(f, "{output}") + } +} + +impl<VAL: ValueType> HeapItem<VAL> { + pub fn new(val: VAL, buk_idx: usize) -> Self { + Self { + val, + map_idx: buk_idx, + } + } +} + +impl<VAL: ValueType> Debug for HeapItem<VAL> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("bucket=")?; + Debug::fmt(&self.map_idx, f)?; + f.write_str(" val=")?; + Debug::fmt(&self.val, f)?; + f.write_str("\n")?; + Ok(()) + } +} + +impl<VAL: ValueType> Eq for HeapItem<VAL> {} + +impl<VAL: ValueType> PartialEq<Self> for HeapItem<VAL> { + fn eq(&self, other: &Self) -> bool { + self.cmp(other) == Ordering::Equal + } +} + +impl<VAL: ValueType> PartialOrd<Self> for HeapItem<VAL> { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + Some(self.cmp(other)) + } +} + +impl<VAL: ValueType> Ord for HeapItem<VAL> { + fn cmp(&self, other: &Self) -> Ordering { + let res = self.val.comp(&other.val); + if res != Ordering::Equal { + return res; + } + self.map_idx.cmp(&other.map_idx) + } +} + +macro_rules! compare_float { + ($($t:ty),+) => { + $(impl Comparable for Option<$t> { + fn comp(&self, other: &Self) -> Ordering { + match (self, other) { + (Some(me), Some(other)) => me.total_cmp(other), + (Some(_), None) => Ordering::Greater, + (None, Some(_)) => Ordering::Less, + (None, None) => Ordering::Equal, + } + } + })+ + + $(impl Comparable for $t { + fn comp(&self, other: &Self) -> Ordering { + self.total_cmp(other) + } + })+ + }; +} + +macro_rules! compare_integer { + ($($t:ty),+) => { + $(impl Comparable for Option<$t> { + fn comp(&self, other: &Self) -> Ordering { + self.cmp(other) + } + })+ + + $(impl Comparable for $t { + fn comp(&self, other: &Self) -> Ordering { + self.cmp(other) + } + })+ + }; +} + +compare_integer!(i8, i16, i32, i64, i128, i256); +compare_integer!(u8, u16, u32, u64); +compare_integer!(IntervalDayTime, IntervalMonthDayNano); +compare_float!(f16, f32, f64); + +/// Returns true if the given data type can be stored in a top-K aggregation heap. +/// +/// Supported types include Arrow primitives (integers, floats, decimals, intervals) +/// and UTF-8 strings (`Utf8`, `LargeUtf8`, `Utf8View`). This is used internally by +/// `PriorityMap::supports()` to validate aggregate value type compatibility. +pub fn is_supported_heap_type(vt: &DataType) -> bool { + vt.is_primitive() + || matches!( + vt, + DataType::Utf8 | DataType::Utf8View | DataType::LargeUtf8 + ) +} + +pub fn new_heap( + limit: usize, + desc: bool, + vt: DataType, +) -> Result<Box<dyn ArrowHeap + Send>> { + if matches!( + vt, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Ok(Box::new(StringHeap::new(limit, desc, vt))); + } + + macro_rules! downcast_helper { + ($vt:ty, $d:ident) => { + return Ok(Box::new(PrimitiveHeap::<$vt>::new(limit, desc, vt))) + }; + } + + downcast_primitive! { + vt => (downcast_helper, vt), + _ => {} + } + + Err(exec_datafusion_err!( + "Unsupported TopK aggregate value type: {vt:?}" + )) +} + +#[cfg(test)] +mod tests { + use insta::assert_snapshot; + + use super::*; + + #[test] + fn should_append() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::new(10, false); + heap.append_or_replace(1, 1, &mut map); + + let actual = heap.to_string(); + assert_snapshot!(actual, @"val=1 idx=0, bucket=1"); + + Ok(()) + } + + #[test] + fn should_heapify_up() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::new(10, false); + + heap.append_or_replace(1, 1, &mut map); + assert_eq!(map, vec![]); + + heap.append_or_replace(2, 2, &mut map); + assert_eq!(map, vec![(2, 0), (1, 1)]); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=2 idx=0, bucket=2 + └── val=1 idx=1, bucket=1 + "); + + Ok(()) + } + + #[test] + fn should_heapify_down() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::new(3, false); + + heap.append_or_replace(1, 1, &mut map); + heap.append_or_replace(2, 2, &mut map); + heap.append_or_replace(3, 3, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=3 idx=0, bucket=3 + ├── val=1 idx=1, bucket=1 + └── val=2 idx=2, bucket=2 + "); + + let mut map = vec![]; + heap.append_or_replace(0, 0, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=2 idx=0, bucket=2 + ├── val=1 idx=1, bucket=1 + └── val=0 idx=2, bucket=0 + "); + assert_eq!(map, vec![(2, 0), (0, 2)]); + + Ok(()) + } + + #[test] + fn should_replace() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::new(4, false); + + heap.append_or_replace(1, 1, &mut map); + heap.append_or_replace(2, 2, &mut map); + heap.append_or_replace(3, 3, &mut map); + heap.append_or_replace(4, 4, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=4 idx=0, bucket=4 + ├── val=3 idx=1, bucket=3 + │ └── val=1 idx=3, bucket=1 + └── val=2 idx=2, bucket=2 + "); + + let mut map = vec![]; + heap.replace_if_better(1, 0, &mut map); + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=4 idx=0, bucket=4 + ├── val=1 idx=1, bucket=1 + │ └── val=0 idx=3, bucket=3 + └── val=2 idx=2, bucket=2 + "); + assert_eq!(map, vec![(1, 1), (3, 3)]); + + Ok(()) + } + + #[test] + fn should_find_worst() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::new(10, false); + + heap.append_or_replace(1, 1, &mut map); + heap.append_or_replace(2, 2, &mut map); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=2 idx=0, bucket=2 + └── val=1 idx=1, bucket=1 + "); + + assert_eq!(heap.worst_val(), Some(&2)); + assert_eq!(heap.worst_map_idx(), 2); + + Ok(()) + } + + #[test] + fn should_drain() -> Result<()> { + let mut map = vec![]; + let mut heap = TopKHeap::new(10, false); + + heap.append_or_replace(1, 1, &mut map); + heap.append_or_replace(2, 2, &mut map); + + let actual = heap.to_string(); + assert_snapshot!(actual, @r" + val=2 idx=0, bucket=2 + └── val=1 idx=1, bucket=1 + "); + + let (vals, map_idxs) = heap.drain(); + assert_eq!(vals, vec![1, 2]); + assert_eq!(map_idxs, vec![1, 2]); + assert_eq!(heap.len(), 0); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/topk/mod.rs b/native/vendor/datafusion-physical-plan/src/aggregates/topk/mod.rs new file mode 100644 index 0000000..c6a0f40 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/topk/mod.rs
@@ -0,0 +1,22 @@ +// 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. + +//! TopK functionality for aggregates + +pub mod hash_table; +pub mod heap; +pub mod priority_map;
diff --git a/native/vendor/datafusion-physical-plan/src/aggregates/topk/priority_map.rs b/native/vendor/datafusion-physical-plan/src/aggregates/topk/priority_map.rs new file mode 100644 index 0000000..f46cb22 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/aggregates/topk/priority_map.rs
@@ -0,0 +1,805 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A `Map<K, V>` / `PriorityQueue` combo that evicts the worst values after reaching `capacity` + +use crate::aggregates::topk::hash_table::{ArrowHashTable, InsertKind, new_hash_table}; +use crate::aggregates::topk::heap::{ArrowHeap, new_heap}; +use arrow::array::{ArrayRef, new_null_array}; +use arrow::compute::concat; +use arrow::datatypes::DataType; +use datafusion_common::Result; + +/// A `Map<K, V>` / `PriorityQueue` combo that evicts the worst values after reaching `capacity` +pub struct PriorityMap { + map: Box<dyn ArrowHashTable + Send>, + heap: Box<dyn ArrowHeap + Send>, + capacity: usize, + mapper: Vec<(usize, usize)>, + val_type: DataType, + /// Mirror of the map's all-NULL group count, kept as a plain field so the + /// per-row `insert` path can check it without a `dyn` call (measured to + /// regress the topk_aggregate benchmarks when read through the trait) + null_count: usize, +} + +impl PriorityMap { + pub fn new( + key_type: DataType, + val_type: DataType, + capacity: usize, + descending: bool, + ) -> Result<Self> { + Ok(Self { + map: new_hash_table(capacity, key_type)?, + heap: new_heap(capacity, descending, val_type.clone())?, + capacity, + mapper: Vec::with_capacity(capacity), + val_type, + null_count: 0, + }) + } + + pub fn set_batch(&mut self, ids: ArrayRef, vals: ArrayRef) { + self.map.set_batch(ids); + self.heap.set_batch(vals); + } + + pub fn insert(&mut self, row_idx: usize) -> Result<()> { + assert!(self.map.len() <= self.capacity, "Overflow"); + debug_assert_eq!(self.null_count, 0); + + // if we're full, and the new val is worse than all our values, just bail + if self.heap.is_worse(row_idx) { + return Ok(()); + } + self.insert_eligible(row_idx) + } + + /// Insert a value while all-NULL groups are being tracked. This is kept + /// separate from [`Self::insert`] so the common no-NULL path does not pay + /// for NULL bookkeeping on every row. + pub fn insert_with_null_groups(&mut self, row_idx: usize) -> Result<()> { + // valued groups are capped at `capacity`; up to `capacity` additional + // all-NULL groups may be tracked alongside them + assert!(self.map.len() <= 2 * self.capacity, "Overflow"); + + if self.heap.is_worse(row_idx) { + // A group that was registered as all-NULL now has a value that + // loses to the current top-k: it can no longer reach the top-k, + // but it must not be emitted with a NULL value either + if self.null_count > 0 && self.map.remove_if_null(row_idx) { + self.null_count -= 1; + } + return Ok(()); + } + self.insert_eligible(row_idx) + } + + fn insert_eligible(&mut self, row_idx: usize) -> Result<()> { + let map = &mut self.mapper; + + // handle new groups we haven't seen yet + map.clear(); + let replace_idx = self.heap.worst_map_idx(); + + let (map_idx, kind) = self.map.find_or_insert(row_idx, replace_idx); + if kind == InsertKind::ReplacedNull { + self.null_count -= 1; + } + if kind != InsertKind::Existing { + self.heap.insert(row_idx, map_idx, map); + self.map.update_heap_idx(map); + return Ok(()); + }; + + // this is a value for an existing group + map.clear(); + let heap_idx = self.map.heap_idx_at(map_idx); + self.heap.replace_if_better(heap_idx, row_idx, map); + self.map.update_heap_idx(map); + + Ok(()) + } + + pub fn has_null_groups(&self) -> bool { + self.null_count > 0 + } + + /// Track a group whose aggregate values are all NULL, so it can be emitted + /// with a NULL value. MIN/MAX ignore NULL inputs, but an all-NULL group + /// must still appear in the aggregation output; such groups all tie on the + /// sort key, so tracking up to `capacity` of them preserves top-k semantics. + pub fn insert_null(&mut self, row_idx: usize) { + assert!(self.map.len() <= 2 * self.capacity, "Overflow"); + if self.map.insert_null(row_idx) { + self.null_count += 1; + } + } + + pub fn emit(&mut self) -> Result<Vec<ArrayRef>> { + let (vals, mut map_idxs) = self.heap.drain(); + // Groups whose values are all NULL are tracked in the map only; + // append them with a NULL value so they are not lost from the output + let null_idxs = self.map.null_map_idxs(); + let vals = if null_idxs.is_empty() { + vals + } else { + map_idxs.extend(null_idxs.iter().copied()); + let nulls = new_null_array(&self.val_type, null_idxs.len()); + concat(&[vals.as_ref(), nulls.as_ref()])? + }; + let ids = self.map.take_all(map_idxs); + self.null_count = 0; + Ok(vec![ids, vals]) + } + + pub fn is_empty(&self) -> bool { + self.map.len() == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + Int64Array, LargeStringArray, RecordBatch, StringArray, StringViewArray, + }; + use arrow::datatypes::{Field, Schema, SchemaRef}; + use arrow::util::pretty::pretty_format_batches; + use insta::assert_snapshot; + use std::sync::Arc; + + #[test] + fn should_append_with_utf8view() -> Result<()> { + let ids: ArrayRef = Arc::new(StringViewArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let mut agg = PriorityMap::new(DataType::Utf8View, DataType::Int64, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema_utf8view(), cols)?; + let batch_schema = batch.schema(); + assert_eq!(batch_schema.fields[0].data_type(), &DataType::Utf8View); + + let actual = format!("{}", pretty_format_batches(&[batch])?); + let expected = r#" ++----------+--------------+ +| trace_id | timestamp_ms | ++----------+--------------+ +| 1 | 1 | ++----------+--------------+ + "# + .trim(); + assert_eq!(actual, expected); + + Ok(()) + } + + #[test] + fn should_append_with_large_utf8() -> Result<()> { + let ids: ArrayRef = Arc::new(LargeStringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let mut agg = PriorityMap::new(DataType::LargeUtf8, DataType::Int64, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_large_schema(), cols)?; + let batch_schema = batch.schema(); + assert_eq!(batch_schema.fields[0].data_type(), &DataType::LargeUtf8); + + let actual = format!("{}", pretty_format_batches(&[batch])?); + let expected = r#" ++----------+--------------+ +| trace_id | timestamp_ms | ++----------+--------------+ +| 1 | 1 | ++----------+--------------+ + "# + .trim(); + assert_eq!(actual, expected); + + Ok(()) + } + + #[test] + fn should_append() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 1 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_ignore_higher_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 1 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_ignore_lower_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["2", "1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![2, 1])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 2 | 2 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_ignore_higher_same_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 1 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_ignore_lower_same_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![2, 1])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 2 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_accept_lower_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["2", "1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![2, 1])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 1 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_accept_higher_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 2 | 2 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_accept_lower_for_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![2, 1])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 1 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_accept_higher_for_group() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 2 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_track_lexicographic_min_utf8_value() -> Result<()> { + let ids: ArrayRef = Arc::new(Int64Array::from(vec![1, 1])); + let vals: ArrayRef = Arc::new(StringArray::from(vec!["zulu", "alpha"])); + let mut agg = PriorityMap::new(DataType::Int64, DataType::Utf8, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema_value(DataType::Utf8), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + + assert_snapshot!(actual, @r#" ++----------+--------------+ +| trace_id | timestamp_ms | ++----------+--------------+ +| 1 | alpha | ++----------+--------------+ + "#); + + Ok(()) + } + + #[test] + fn should_track_lexicographic_max_utf8_value_desc() -> Result<()> { + let ids: ArrayRef = Arc::new(Int64Array::from(vec![1, 1])); + let vals: ArrayRef = Arc::new(StringArray::from(vec!["alpha", "zulu"])); + let mut agg = PriorityMap::new(DataType::Int64, DataType::Utf8, 1, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema_value(DataType::Utf8), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + + assert_snapshot!(actual, @r#" ++----------+--------------+ +| trace_id | timestamp_ms | ++----------+--------------+ +| 1 | zulu | ++----------+--------------+ + "#); + + Ok(()) + } + + #[test] + fn should_track_large_utf8_values() -> Result<()> { + let ids: ArrayRef = Arc::new(Int64Array::from(vec![1, 1])); + let vals: ArrayRef = Arc::new(LargeStringArray::from(vec!["zulu", "alpha"])); + let mut agg = PriorityMap::new(DataType::Int64, DataType::LargeUtf8, 1, false)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema_value(DataType::LargeUtf8), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + + assert_snapshot!(actual, @r#" ++----------+--------------+ +| trace_id | timestamp_ms | ++----------+--------------+ +| 1 | alpha | ++----------+--------------+ + "#); + + Ok(()) + } + + #[test] + fn should_track_utf8_view_values() -> Result<()> { + let ids: ArrayRef = Arc::new(Int64Array::from(vec![1, 1])); + let vals: ArrayRef = Arc::new(StringViewArray::from(vec!["alpha", "zulu"])); + let mut agg = PriorityMap::new(DataType::Int64, DataType::Utf8View, 1, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema_value(DataType::Utf8View), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + + assert_snapshot!(actual, @r#" ++----------+--------------+ +| trace_id | timestamp_ms | ++----------+--------------+ +| 1 | zulu | ++----------+--------------+ + "#); + + Ok(()) + } + + #[test] + fn should_handle_null_ids() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec![Some("1"), None, None])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert(1)?; + agg.insert(2)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | | 3 | + | 1 | 1 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_emit_all_null_groups() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None, None])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + agg.set_batch(ids, vals); + agg.insert_null(0); + agg.insert_null(1); + // re-registering an existing NULL group is a no-op + agg.insert_null(0); + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | | + | 2 | | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_emit_null_groups_alongside_valued_groups() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![Some(7), None, Some(3)])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 3, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert_null(1); + agg.insert_with_null_groups(2)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 7 | + | 3 | 3 | + | 2 | | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_cap_null_groups_at_limit() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3", "4", "5"])); + let vals: ArrayRef = + Arc::new(Int64Array::from(vec![None, None, None, None, None])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, false)?; + agg.set_batch(ids, vals); + for row_idx in 0..5 { + agg.insert_null(row_idx); + } + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | | + | 2 | | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_convert_null_group_to_valued() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + + // group "1" only produces NULLs in the first batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + // group "1" produces a value in a later batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![5])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 5 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_not_duplicate_valued_group_as_null() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, false)?; + + // group "1" produces a value in the first batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![5])); + agg.set_batch(ids, vals); + agg.insert(0)?; + + // group "1" only produces NULLs in a later batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 5 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_evict_worst_when_converting_null_group() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, true)?; + + // group "2" holds the single top-k slot + let ids: ArrayRef = Arc::new(StringArray::from(vec!["2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![10])); + agg.set_batch(ids, vals); + agg.insert(0)?; + + // group "1" starts out all-NULL + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + // group "1" produces a better value and evicts group "2" + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![20])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 20 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_drop_null_group_that_loses_to_topk() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, true)?; + + // group "1" starts out all-NULL + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + // group "2" fills the single top-k slot + let ids: ArrayRef = Arc::new(StringArray::from(vec!["2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![10])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + // group "1" produces a value that loses to the current top-k: the + // group can no longer reach the top-k and must not be emitted as NULL + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![5])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 2 | 10 | + +----------+--------------+ + " + ); + + Ok(()) + } + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("trace_id", DataType::Utf8, true), + Field::new("timestamp_ms", DataType::Int64, true), + ])) + } + + fn test_schema_utf8view() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("trace_id", DataType::Utf8View, true), + Field::new("timestamp_ms", DataType::Int64, true), + ])) + } + + fn test_large_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("trace_id", DataType::LargeUtf8, true), + Field::new("timestamp_ms", DataType::Int64, true), + ])) + } + + fn test_schema_value(value_type: DataType) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("trace_id", DataType::Int64, true), + Field::new("timestamp_ms", value_type, true), + ])) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/analyze.rs b/native/vendor/datafusion-physical-plan/src/analyze.rs new file mode 100644 index 0000000..d151982 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/analyze.rs
@@ -0,0 +1,566 @@ +// 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. + +//! Defines the ANALYZE operator + +use std::sync::Arc; + +use super::stream::{RecordBatchReceiverStream, RecordBatchStreamAdapter}; +use super::{ + DisplayAs, Distribution, ExecutionPlanProperties, PlanProperties, + SendableRecordBatchStream, +}; +use crate::display::DisplayableExecutionPlan; +use crate::execution_plan::EvaluationType; +use crate::metrics::{MetricCategory, MetricType}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, +}; + +use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; +use datafusion_common::format::ExplainFormat; +use datafusion_common::instant::Instant; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + DataFusionError, Result, assert_eq_or_internal_err, internal_err, +}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::PhysicalExpr; + +use futures::StreamExt; + +/// `EXPLAIN ANALYZE` execution plan operator. This operator runs its input, +/// discards the results, and then prints out an annotated plan with metrics +#[derive(Debug, Clone)] +pub struct AnalyzeExec { + /// Control how much extra to print + verbose: bool, + /// If statistics should be displayed + show_statistics: bool, + /// Which metric categories should be displayed + metric_types: Vec<MetricType>, + /// Optional filter by semantic category (rows / bytes / timing). + metric_categories: Option<Vec<MetricCategory>>, + /// Output format for the rendered plan + metrics. + format: ExplainFormat, + /// The input plan (the plan being analyzed) + pub(crate) input: Arc<dyn ExecutionPlan>, + /// The output schema for RecordBatches of this exec node + schema: SchemaRef, + cache: Arc<PlanProperties>, +} + +/// Builder for [`AnalyzeExec`]. +/// +/// Builder for [AnalyzeExec]. +pub struct AnalyzeExecBuilder { + verbose: bool, + show_statistics: bool, + input: Arc<dyn ExecutionPlan>, + schema: SchemaRef, + metric_types: Vec<MetricType>, + metric_categories: Option<Vec<MetricCategory>>, + format: ExplainFormat, +} + +impl AnalyzeExecBuilder { + pub fn new( + verbose: bool, + show_statistics: bool, + input: Arc<dyn ExecutionPlan>, + schema: SchemaRef, + ) -> Self { + Self { + verbose, + show_statistics, + input, + schema, + metric_types: vec![MetricType::Summary, MetricType::Dev], + metric_categories: None, + format: ExplainFormat::Indent, + } + } + + pub fn with_metric_types(mut self, metric_types: Vec<MetricType>) -> Self { + self.metric_types = metric_types; + self + } + + pub fn with_metric_categories( + mut self, + metric_categories: Option<Vec<MetricCategory>>, + ) -> Self { + self.metric_categories = metric_categories; + self + } + + pub fn with_format(mut self, format: ExplainFormat) -> Self { + self.format = format; + self + } + + pub fn build(self) -> AnalyzeExec { + let cache = + AnalyzeExec::compute_properties(&self.input, Arc::clone(&self.schema)); + AnalyzeExec { + verbose: self.verbose, + show_statistics: self.show_statistics, + metric_types: self.metric_types, + metric_categories: self.metric_categories, + format: self.format, + input: self.input, + schema: self.schema, + cache: Arc::new(cache), + } + } +} + +impl AnalyzeExec { + /// Returns a builder for constructing an [`AnalyzeExec`]. + pub fn builder( + verbose: bool, + show_statistics: bool, + input: Arc<dyn ExecutionPlan>, + schema: SchemaRef, + ) -> AnalyzeExecBuilder { + AnalyzeExecBuilder::new(verbose, show_statistics, input, schema) + } + + /// Access to verbose + pub fn verbose(&self) -> bool { + self.verbose + } + + /// Access to show_statistics + pub fn show_statistics(&self) -> bool { + self.show_statistics + } + + /// Access to metric_categories + pub fn metric_categories(&self) -> Option<&[MetricCategory]> { + self.metric_categories.as_deref() + } + + /// Access to format + pub fn format(&self) -> &ExplainFormat { + &self.format + } + + /// The input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + ) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + input.pipeline_behavior(), + input.boundedness(), + ) + .with_evaluation_type(EvaluationType::Eager) + } +} + +impl DisplayAs for AnalyzeExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "AnalyzeExec verbose={}", self.verbose) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for AnalyzeExec { + fn name(&self) -> &'static str { + "AnalyzeExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new( + AnalyzeExec::builder( + self.verbose, + self.show_statistics, + children.pop().unwrap(), + Arc::clone(&self.schema), + ) + .with_metric_types(self.metric_types.clone()) + .with_metric_categories(self.metric_categories.clone()) + .with_format(self.format.clone()) + .build(), + )) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + assert_eq_or_internal_err!( + partition, + 0, + "AnalyzeExec invalid partition. Expected 0, got {partition}" + ); + + // Gather futures that will run each input partition in + // parallel (on a separate tokio task) using a JoinSet to + // cancel outstanding futures on drop + let num_input_partitions = self.input.output_partitioning().partition_count(); + let mut builder = + RecordBatchReceiverStream::builder(self.schema(), num_input_partitions); + + for input_partition in 0..num_input_partitions { + builder.run_input( + Arc::clone(&self.input), + input_partition, + Arc::clone(&context), + ); + } + + // Create future that computes the final output + let start = Instant::now(); + let captured_input = Arc::clone(&self.input); + let captured_schema = Arc::clone(&self.schema); + let verbose = self.verbose; + let show_statistics = self.show_statistics; + let metric_types = self.metric_types.clone(); + let metric_categories = self.metric_categories.clone(); + let format = self.format.clone(); + + // future that gathers the results from all the tasks in the + // JoinSet that computes the overall row count and final + // record batch + let mut input_stream = builder.build(); + let output = async move { + let mut total_rows = 0; + while let Some(batch) = input_stream.next().await.transpose()? { + total_rows += batch.num_rows(); + } + drop(input_stream); + + let duration = Instant::now() - start; + create_output_batch( + verbose, + show_statistics, + total_rows, + duration, + &captured_input, + &captured_schema, + &metric_types, + metric_categories.as_deref(), + &format, + ) + }; + + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + futures::stream::once(output), + ))) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `AnalyzeExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + verbose, + show_statistics, + // TODO: not on the wire. `AnalyzeExecBuilder` always resets this to + // `[Summary, Dev]`, so a non-default selection is lost on + // round-trip. Fixing it needs a new proto field. + metric_types: _, + metric_categories, + format, + input, + schema, + // Derived at construction from `input` and `schema`. + cache: _, + } = self; + + let input = ctx.encode_child(input)?; + let (has_metric_categories, metric_categories) = match metric_categories { + Some(categories) => { + (true, categories.iter().map(ToString::to_string).collect()) + } + None => (false, vec![]), + }; + let format = match format { + ExplainFormat::Indent => protobuf::ExplainFormat::Indent, + ExplainFormat::Tree => protobuf::ExplainFormat::Tree, + ExplainFormat::PostgresJSON => protobuf::ExplainFormat::Pgjson, + ExplainFormat::Graphviz => protobuf::ExplainFormat::Graphviz, + } as i32; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Analyze(Box::new( + protobuf::AnalyzeExecNode { + verbose: *verbose, + show_statistics: *show_statistics, + input: Some(Box::new(input)), + schema: Some(schema.as_ref().try_into()?), + has_metric_categories, + metric_categories, + format, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AnalyzeExec { + /// Reconstruct an [`AnalyzeExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let analyze = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Analyze, + "AnalyzeExec", + ); + // Exhaustive destructure: a new field on `AnalyzeExecNode` is a compile + // error here rather than a silently ignored wire field. + let protobuf::AnalyzeExecNode { + verbose, + show_statistics, + input, + schema, + has_metric_categories, + metric_categories, + format, + } = analyze.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "AnalyzeExec", "input")?; + let metric_categories = if *has_metric_categories { + Some( + metric_categories + .iter() + .map(|category| category.parse::<MetricCategory>()) + .collect::<Result<Vec<_>>>()?, + ) + } else { + None + }; + let proto_format = protobuf::ExplainFormat::try_from(*format).map_err(|_| { + DataFusionError::Internal(format!( + "Received an AnalyzeExecNode message with unknown ExplainFormat {format}" + )) + })?; + let format = match proto_format { + protobuf::ExplainFormat::Indent => ExplainFormat::Indent, + protobuf::ExplainFormat::Tree => ExplainFormat::Tree, + protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON, + protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz, + }; + let schema = schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "AnalyzeExec is missing required field 'schema'" + ) + })?; + Ok(Arc::new( + AnalyzeExec::builder( + *verbose, + *show_statistics, + input, + Arc::new(arrow::datatypes::Schema::try_from(schema)?), + ) + .with_metric_categories(metric_categories) + .with_format(format) + .build(), + )) + } +} + +/// Creates the output of AnalyzeExec as a RecordBatch +#[expect(clippy::too_many_arguments)] +fn create_output_batch( + verbose: bool, + show_statistics: bool, + total_rows: usize, + duration: std::time::Duration, + input: &Arc<dyn ExecutionPlan>, + schema: &SchemaRef, + metric_types: &[MetricType], + metric_categories: Option<&[MetricCategory]>, + format: &ExplainFormat, +) -> Result<RecordBatch> { + let mut type_builder = StringBuilder::with_capacity(1, 1024); + let mut plan_builder = StringBuilder::with_capacity(1, 1024); + + match format { + ExplainFormat::Indent => { + // TODO use some sort of enum rather than strings? + type_builder.append_value("Plan with Metrics"); + let annotated_plan = DisplayableExecutionPlan::with_metrics(input.as_ref()) + .set_metric_types(metric_types.to_vec()) + .set_metric_categories(metric_categories.map(|c| c.to_vec())) + .set_show_statistics(show_statistics) + .indent(verbose) + .to_string(); + plan_builder.append_value(annotated_plan); + // Verbose output + // TODO make this more sophisticated + if verbose { + type_builder.append_value("Plan with Full Metrics"); + let annotated_plan = + DisplayableExecutionPlan::with_full_metrics(input.as_ref()) + .set_metric_types(metric_types.to_vec()) + .set_metric_categories(metric_categories.map(|c| c.to_vec())) + .set_show_statistics(show_statistics) + .indent(verbose) + .to_string(); + plan_builder.append_value(annotated_plan); + type_builder.append_value("Output Rows"); + plan_builder.append_value(total_rows.to_string()); + type_builder.append_value("Duration"); + plan_builder.append_value(format!("{duration:?}")); + } + } + ExplainFormat::PostgresJSON => { + // `show_statistics` is intentionally not forwarded here: the pgjson + // renderer does not emit statistics, and the planner rejects the + // `show_statistics` + pgjson combination up front. + type_builder.append_value("Plan with Metrics"); + let mut displayable = if verbose { + DisplayableExecutionPlan::with_full_metrics(input.as_ref()) + } else { + DisplayableExecutionPlan::with_metrics(input.as_ref()) + }; + displayable = displayable + .set_metric_types(metric_types.to_vec()) + .set_metric_categories(metric_categories.map(|c| c.to_vec())); + if verbose { + displayable = displayable.set_summary(Some(total_rows), Some(duration)); + } + plan_builder.append_value(displayable.pgjson(verbose).to_string()); + } + ExplainFormat::Tree | ExplainFormat::Graphviz => { + return internal_err!("AnalyzeExec does not support {format} output format"); + } + } + + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(type_builder.finish()), + Arc::new(plan_builder.finish()), + ], + ) + .map_err(DataFusionError::from) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + collect, + test::{ + assert_is_pending, + exec::{BlockingExec, assert_strong_count_converges_to_zero}, + }, + }; + + use arrow::datatypes::{DataType, Field, Schema}; + use futures::FutureExt; + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let analyze_exec = + Arc::new(AnalyzeExec::builder(true, false, blocking_exec, schema).build()); + + let fut = collect(analyze_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/async_func.rs b/native/vendor/datafusion-physical-plan/src/async_func.rs new file mode 100644 index 0000000..f3ef13d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/async_func.rs
@@ -0,0 +1,550 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::coalesce::LimitedBatchCoalescer; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, + validate_child_count, +}; +use arrow::array::RecordBatch; +use arrow_schema::{FieldRef, Fields, Schema, SchemaRef}; +use datafusion_common::Result; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::ScalarFunctionExpr; +use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::metrics::{BaselineMetrics, RecordOutput}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use futures::Stream; +use futures::stream::StreamExt; +use log::trace; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, ready}; + +/// This structure evaluates a set of async expressions on a record +/// batch producing a new record batch +/// +/// The schema of the output of the AsyncFuncExec is: +/// Input columns followed by one column for each async expression +#[derive(Debug, Clone)] +pub struct AsyncFuncExec { + /// The async expressions to evaluate + async_exprs: Vec<Arc<AsyncFuncExpr>>, + input: Arc<dyn ExecutionPlan>, + cache: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl AsyncFuncExec { + pub fn try_new( + async_exprs: Vec<Arc<AsyncFuncExpr>>, + input: Arc<dyn ExecutionPlan>, + ) -> Result<Self> { + let async_fields = async_exprs + .iter() + .map(|async_expr| async_expr.return_field(input.schema().as_ref())) + .collect::<Result<Vec<FieldRef>>>()?; + + // compute the output schema: input schema then async expressions + let fields: Fields = input + .schema() + .fields() + .iter() + .cloned() + .chain(async_fields) + .collect(); + + let schema = Arc::new(Schema::new(fields)); + let tuples = async_exprs + .iter() + .map(|expr| (Arc::clone(&expr.func), expr.name().to_string())) + .collect::<Vec<_>>(); + let async_expr_mapping = ProjectionMapping::try_new(tuples, &input.schema())?; + let cache = + AsyncFuncExec::compute_properties(&input, schema, &async_expr_mapping)?; + Ok(Self { + input, + async_exprs, + cache: Arc::new(cache), + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + /// This function creates the cache object that stores the plan properties + /// such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + async_expr_mapping: &ProjectionMapping, + ) -> Result<PlanProperties> { + Ok(PlanProperties::new( + input + .equivalence_properties() + .project(async_expr_mapping, schema), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )) + } + + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `AsyncFuncExec` serializes itself via `AsyncFuncExec::try_to_proto`, which reads the field directly. There is no replacement; please open an issue if you have a use case for it." + )] + pub fn async_exprs(&self) -> &[Arc<AsyncFuncExpr>] { + &self.async_exprs + } + + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } +} + +impl DisplayAs for AsyncFuncExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + let expr: Vec<String> = self + .async_exprs + .iter() + .map(|async_expr| async_expr.to_string()) + .collect(); + let exprs = expr.join(", "); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "AsyncFuncExec: async_expr=[{exprs}]") + } + DisplayFormatType::TreeRender => { + writeln!(f, "format=async_expr")?; + writeln!(f, "async_expr={exprs}")?; + Ok(()) + } + } + } +} + +impl ExecutionPlan for AsyncFuncExec { + fn name(&self) -> &str { + "async_func" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + crate::apply_expression_roots( + self.async_exprs + .iter() + .cloned() + .map(|expr| expr as Arc<dyn PhysicalExpr>), + f, + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(AsyncFuncExec::try_new( + self.async_exprs.clone(), + children.swap_remove(0), + )?)), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start AsyncFuncExpr::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + // first execute the input stream + let input_stream = self.input.execute(partition, Arc::clone(&context))?; + + // TODO: Track `elapsed_compute` in `BaselineMetrics` + // Issue: <https://github.com/apache/datafusion/issues/19658> + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + + // now, for each record batch, evaluate the async expressions and add the columns to the result + let async_exprs_captured = Arc::new(self.async_exprs.clone()); + let schema_captured = self.schema(); + let config_options_ref = Arc::clone(context.session_config().options()); + + let coalesced_input_stream = CoalesceInputStream { + input_stream, + batch_coalescer: LimitedBatchCoalescer::new( + Arc::clone(&self.input.schema()), + config_options_ref.execution.batch_size.get(), + None, + ), + }; + + let stream_with_async_functions = coalesced_input_stream.then(move |batch| { + // need to clone *again* to capture the async_exprs and schema in the + // stream and satisfy lifetime requirements. + let async_exprs_captured = Arc::clone(&async_exprs_captured); + let schema_captured = Arc::clone(&schema_captured); + let config_options = Arc::clone(&config_options_ref); + let baseline_metrics_captured = baseline_metrics.clone(); + + async move { + let batch = batch?; + // append the result of evaluating the async expressions to the output + let mut output_arrays = batch.columns().to_vec(); + for async_expr in async_exprs_captured.iter() { + let output = async_expr + .invoke_with_args(&batch, Arc::clone(&config_options)) + .await?; + output_arrays.push(output.to_array(batch.num_rows())?); + } + let batch = RecordBatch::try_new(schema_captured, output_arrays)?; + + Ok(batch.record_output(&baseline_metrics_captured)) + } + }); + + // Adapt the stream with the output schema + let adapter = + RecordBatchStreamAdapter::new(self.schema(), stream_with_async_functions); + Ok(Box::pin(adapter)) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `AsyncFuncExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + async_exprs, + input, + // Derived at construction by `AsyncFuncExec::compute_properties`. + cache: _, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + } = self; + + let input = ctx.encode_child(input)?; + let async_expr_names = async_exprs.iter().map(|e| e.name().to_string()).collect(); + let async_exprs = ctx.encode_expressions(async_exprs.iter().map(|e| &e.func))?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc(Box::new( + protobuf::AsyncFuncExecNode { + input: Some(Box::new(input)), + async_exprs, + async_expr_names, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AsyncFuncExec { + /// Reconstruct an [`AsyncFuncExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. Child plans and expressions are decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_common::assert_eq_or_internal_err; + use datafusion_proto_models::protobuf; + let async_func = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc, + "AsyncFuncExec", + ); + // Exhaustive destructure: a new field on `AsyncFuncExecNode` is a + // compile error here rather than a silently ignored wire field. + let protobuf::AsyncFuncExecNode { + input, + async_exprs, + async_expr_names, + } = async_func.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "AsyncFuncExec", "input")?; + let input_schema = input.schema(); + assert_eq_or_internal_err!( + async_exprs.len(), + async_expr_names.len(), + "AsyncFuncExecNode async_exprs length does not match async_expr_names" + ); + let async_exprs = async_exprs + .iter() + .zip(async_expr_names.iter()) + .map(|(expr, name)| { + let physical_expr = ctx.decode_expr(expr, input_schema.as_ref())?; + Ok(Arc::new(AsyncFuncExpr::try_new( + name.clone(), + physical_expr, + input_schema.as_ref(), + )?)) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) + } +} + +struct CoalesceInputStream { + input_stream: Pin<Box<dyn RecordBatchStream + Send>>, + batch_coalescer: LimitedBatchCoalescer, +} + +impl Stream for CoalesceInputStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let mut completed = false; + + loop { + if let Some(batch) = self.batch_coalescer.next_completed_batch() { + return Poll::Ready(Some(Ok(batch))); + } + + if completed { + return Poll::Ready(None); + } + + match ready!(self.input_stream.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if let Err(err) = self.batch_coalescer.push_batch(batch) { + return Poll::Ready(Some(Err(err))); + } + } + Some(err) => { + return Poll::Ready(Some(err)); + } + None => { + completed = true; + // Release the input pipeline's resources. + let input_schema = self.input_stream.schema(); + self.input_stream = + Box::pin(EmptyRecordBatchStream::new(input_schema)); + if let Err(err) = self.batch_coalescer.finish() { + return Poll::Ready(Some(Err(err))); + } + } + } + } + } +} + +const ASYNC_FN_PREFIX: &str = "__async_fn_"; + +/// Maps async_expressions to new columns +/// +/// The output of the async functions are appended, in order, to the end of the input schema +#[derive(Debug)] +pub struct AsyncMapper { + /// the number of columns in the input plan + /// used to generate the output column names. + /// the first async expr is `__async_fn_0`, the second is `__async_fn_1`, etc + num_input_columns: usize, + /// the expressions to map + pub async_exprs: Vec<Arc<AsyncFuncExpr>>, +} + +impl AsyncMapper { + pub fn new(num_input_columns: usize) -> Self { + Self { + num_input_columns, + async_exprs: Vec::new(), + } + } + + pub fn is_empty(&self) -> bool { + self.async_exprs.is_empty() + } + + pub fn next_column_name(&self) -> String { + format!("{}{}", ASYNC_FN_PREFIX, self.async_exprs.len()) + } + + /// Finds any references to async functions in the expression and adds them to the map + pub fn find_references( + &mut self, + physical_expr: &Arc<dyn PhysicalExpr>, + schema: &Schema, + ) -> Result<()> { + // recursively look for references to async functions + physical_expr.apply(|expr| { + if let Some(scalar_func_expr) = expr.downcast_ref::<ScalarFunctionExpr>() + && scalar_func_expr.fun().as_async().is_some() + { + let next_name = self.next_column_name(); + self.async_exprs.push(Arc::new(AsyncFuncExpr::try_new( + next_name, + Arc::clone(expr), + schema, + )?)); + } + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) + } + + /// If the expression matches any of the async functions, return the new column + pub fn map_expr( + &self, + expr: Arc<dyn PhysicalExpr>, + ) -> Transformed<Arc<dyn PhysicalExpr>> { + // find the first matching async function if any + let Some(idx) = + self.async_exprs + .iter() + .enumerate() + .find_map(|(idx, async_expr)| { + if async_expr.func == Arc::clone(&expr) { + Some(idx) + } else { + None + } + }) + else { + return Transformed::no(expr); + }; + // rewrite in terms of the output column + Transformed::yes(self.output_column(idx)) + } + + /// return the output column for the async function at index idx + pub fn output_column(&self, idx: usize) -> Arc<dyn PhysicalExpr> { + let async_expr = &self.async_exprs[idx]; + let output_idx = self.num_input_columns + idx; + Arc::new(Column::new(async_expr.name(), output_idx)) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{RecordBatch, UInt32Array}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_execution::{TaskContext, config::SessionConfig}; + use futures::StreamExt; + + use crate::{ExecutionPlan, async_func::AsyncFuncExec, test::TestMemoryExec}; + + #[tokio::test] + async fn test_async_fn_with_coalescing() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6]))], + )?; + + let batches: Vec<RecordBatch> = std::iter::repeat_n(batch, 50).collect(); + + let session_config = SessionConfig::new().with_batch_size(200); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + let test_exec = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let exec = AsyncFuncExec::try_new(vec![], test_exec)?; + + let mut stream = exec.execute(0, Arc::clone(&task_ctx))?; + let batch = stream + .next() + .await + .expect("expected to get a record batch")?; + assert_eq!(200, batch.num_rows()); + let batch = stream + .next() + .await + .expect("expected to get a record batch")?; + assert_eq!(100, batch.num_rows()); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/buffer.rs b/native/vendor/datafusion-physical-plan/src/buffer.rs new file mode 100644 index 0000000..24cca6b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/buffer.rs
@@ -0,0 +1,753 @@ +// 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. + +//! [`BufferExec`] decouples production and consumption on messages by buffering the input in the +//! background up to a certain capacity. + +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; +use crate::filter_pushdown::{ + ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, SortOrderPushdownResult, validate_child_count, +}; +use arrow::array::RecordBatch; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, Statistics, internal_err}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet, +}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; +use pin_project_lite::pin_project; +use std::fmt; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use tokio::sync::mpsc::UnboundedReceiver; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +/// WARNING: EXPERIMENTAL +/// +/// Decouples production and consumption of record batches with an internal queue per partition, +/// eagerly filling up the capacity of the queues even before any message is requested. +/// +/// ```text +/// ┌───────────────────────────┐ +/// │ BufferExec │ +/// │ │ +/// │┌────── Partition 0 ──────┐│ +/// ││ ┌────┐ ┌────┐││ ┌────┐ +/// ──background poll────────▶│ │ │ ├┼┼───────▶ │ +/// ││ └────┘ └────┘││ └────┘ +/// │└─────────────────────────┘│ +/// │┌────── Partition 1 ──────┐│ +/// ││ ┌────┐ ┌────┐ ┌────┐││ ┌────┐ +/// ──background poll─▶│ │ │ │ │ ├┼┼───────▶ │ +/// ││ └────┘ └────┘ └────┘││ └────┘ +/// │└─────────────────────────┘│ +/// │ │ +/// │ ... │ +/// │ │ +/// │┌────── Partition N ──────┐│ +/// ││ ┌────┐││ ┌────┐ +/// ──background poll───────────────▶│ ├┼┼───────▶ │ +/// ││ └────┘││ └────┘ +/// │└─────────────────────────┘│ +/// └───────────────────────────┘ +/// ``` +/// +/// The capacity is provided in bytes, and for each buffered record batch it will take into account +/// the size reported by [RecordBatch::get_array_memory_size]. +/// +/// If a single record batch exceeds the maximum capacity set in the `capacity` argument, it's still +/// allowed to pass in order to not deadlock the buffer. +/// +/// This is useful for operators that conditionally start polling one of their children only after +/// other child has finished, allowing to perform some early work and accumulating batches in +/// memory so that they can be served immediately when requested. +#[derive(Debug, Clone)] +pub struct BufferExec { + input: Arc<dyn ExecutionPlan>, + properties: Arc<PlanProperties>, + capacity: usize, + metrics: ExecutionPlanMetricsSet, +} + +impl BufferExec { + /// Builds a new [BufferExec] with the provided capacity in bytes. + pub fn new(input: Arc<dyn ExecutionPlan>, capacity: usize) -> Self { + let properties = PlanProperties::clone(input.properties()) + .with_scheduling_type(SchedulingType::Cooperative) + .with_evaluation_type(EvaluationType::Eager); + + Self { + input, + properties: Arc::new(properties), + capacity, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + /// Returns the input [ExecutionPlan] of this [BufferExec]. + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Returns the per-partition capacity in bytes for this [BufferExec]. + pub fn capacity(&self) -> usize { + self.capacity + } +} + +impl DisplayAs for BufferExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "BufferExec: capacity={}", self.capacity) + } + DisplayFormatType::TreeRender => { + writeln!(f, "target_batch_size={}", self.capacity) + } + } + } +} + +impl ExecutionPlan for BufferExec { + fn name(&self) -> &str { + "BufferExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.properties + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let mem_reservation = MemoryConsumer::new(format!("BufferExec[{partition}]")) + .register(context.memory_pool()); + let in_stream = self.input.execute(partition, context)?; + + // Set up the metrics for the stream. + let curr_mem_in = Arc::new(AtomicUsize::new(0)); + let curr_mem_out = Arc::clone(&curr_mem_in); + let mut max_mem_in = 0; + let max_mem = MetricBuilder::new(&self.metrics) + .peak_memory_usage("max_mem_used", partition); + + let curr_queued_in = Arc::new(AtomicUsize::new(0)); + let curr_queued_out = Arc::clone(&curr_queued_in); + let mut max_queued_in = 0; + let max_queued = MetricBuilder::new(&self.metrics) + .with_category(MetricCategory::Rows) + .gauge("max_queued", partition); + + // Capture metrics when an element is queued on the stream. + let in_stream = in_stream.inspect_ok(move |v| { + let size = v.get_array_memory_size(); + let curr_size = curr_mem_in.fetch_add(size, Ordering::Relaxed) + size; + if curr_size > max_mem_in { + max_mem_in = curr_size; + max_mem.set(max_mem_in); + } + + let curr_queued = curr_queued_in.fetch_add(1, Ordering::Relaxed) + 1; + if curr_queued > max_queued_in { + max_queued_in = curr_queued; + max_queued.set(max_queued_in); + } + }); + // Buffer the input. + let out_stream = + MemoryBufferedStream::new(in_stream, self.capacity, mem_reservation); + // Update in the metrics that when an element gets out, some memory gets freed. + let out_stream = out_stream.inspect_ok(move |v| { + curr_mem_out.fetch_sub(v.get_array_memory_size(), Ordering::Relaxed); + curr_queued_out.fetch_sub(1, Ordering::Relaxed); + }); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + out_stream, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::clone(&input_stats[0])) + } + + fn supports_limit_pushdown(&self) -> bool { + self.input.supports_limit_pushdown() + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + match self.input.try_swapping_with_projection(projection)? { + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), + None => Ok(None), + } + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + FilterDescription::from_children(parent_filters, &self.children()) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + Ok(FilterPushdownPropagation::if_all(child_pushdown_result)) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + // CoalesceBatchesExec is transparent for sort ordering - it preserves order + // Delegate to the child and wrap with a new CoalesceBatchesExec + self.input.try_pushdown_sort(order)?.try_map(|new_input| { + Ok(Arc::new(Self::new(new_input, self.capacity)) as Arc<dyn ExecutionPlan>) + }) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Buffer(Box::new( + protobuf::BufferExecNode { + input: Some(Box::new(input)), + capacity: self.capacity() as u64, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl BufferExec { + /// Reconstruct a [`BufferExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let buffer = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Buffer, + "BufferExec", + ); + let input = + ctx.decode_required_child(buffer.input.as_deref(), "BufferExec", "input")?; + Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) + } +} + +/// Represents anything that occupies a capacity in a [MemoryBufferedStream]. +pub trait SizedMessage { + fn size(&self) -> usize; +} + +impl SizedMessage for RecordBatch { + fn size(&self) -> usize { + self.get_array_memory_size() + } +} + +pin_project! { +/// Decouples production and consumption of messages in a stream with an internal queue, eagerly +/// filling it up to the specified maximum capacity even before any message is requested. +/// +/// Allows each message to have a different size, which is taken into account for determining if +/// the queue is full or not. +pub struct MemoryBufferedStream<T: SizedMessage> { + task: SpawnedTask<()>, + batch_rx: UnboundedReceiver<Result<(T, OwnedSemaphorePermit)>>, + memory_reservation: Arc<MemoryReservation>, +}} + +impl<T: Send + SizedMessage + 'static> MemoryBufferedStream<T> { + /// Builds a new [MemoryBufferedStream] with the provided capacity and event handler. + /// + /// This immediately spawns a Tokio task that will start consumption of the input stream. + pub fn new( + mut input: impl Stream<Item = Result<T>> + Unpin + Send + 'static, + capacity: usize, + memory_reservation: MemoryReservation, + ) -> Self { + let semaphore = Arc::new(Semaphore::new(capacity)); + let (batch_tx, batch_rx) = tokio::sync::mpsc::unbounded_channel(); + + let memory_reservation = Arc::new(memory_reservation); + let memory_reservation_clone = Arc::clone(&memory_reservation); + let task = SpawnedTask::spawn(async move { + loop { + // Select on both the input stream and the channel being closed. + // By down this, we abort polling the input as soon as the consumer channel is + // closed. Otherwise, we would need to wait for a full new message to be available + // in order to consider aborting the stream + let item_or_err = tokio::select! { + biased; + _ = batch_tx.closed() => break, + // Catch a panic in the input poll so it surfaces as a stream error + // instead of dropping `batch_tx` and looking like a clean EOF. + polled = AssertUnwindSafe(input.next()).catch_unwind() => { + match polled { + Ok(Some(item_or_err)) => item_or_err, + Ok(None) => break, // stream finished + Err(panic) => { + let msg = panic + .downcast_ref::<&str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::<String>().cloned()) + .unwrap_or_else(|| "unknown panic".to_string()); + let _ = batch_tx.send(internal_err!( + "BufferExec input stream panicked: {msg}" + )); + break; + } + } + } + }; + + let item = match item_or_err { + Ok(batch) => batch, + Err(err) => { + let _ = batch_tx.send(Err(err)); // If there's an error it means the channel was closed, which is fine. + break; + } + }; + + let size = item.size(); + if let Err(err) = memory_reservation.try_grow(size) { + let _ = batch_tx.send(Err(err)); // If there's an error it means the channel was closed, which is fine. + break; + } + + // We need to cap the minimum between amount of permits and the actual size of the + // message. If at any point we try to acquire more permits than the capacity of the + // semaphore, the stream will deadlock. + let capped_size = size.min(capacity) as u32; + + let semaphore = Arc::clone(&semaphore); + let Ok(permit) = semaphore.acquire_many_owned(capped_size).await else { + let _ = batch_tx.send(internal_err!("Closed semaphore in MemoryBufferedStream. This is a bug in DataFusion, please report it!")); + break; + }; + + if batch_tx.send(Ok((item, permit))).is_err() { + break; // stream was closed + }; + } + }); + + Self { + task, + batch_rx, + memory_reservation: memory_reservation_clone, + } + } + + /// Returns the number of queued messages. + pub fn messages_queued(&self) -> usize { + self.batch_rx.len() + } +} + +impl<T: SizedMessage> Stream for MemoryBufferedStream<T> { + type Item = Result<T>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let self_project = self.project(); + match self_project.batch_rx.poll_recv(cx) { + Poll::Ready(Some(Ok((item, _semaphore_permit)))) => { + self_project.memory_reservation.shrink(item.size()); + Poll::Ready(Some(Ok(item))) + } + Poll::Ready(Some(Err(err))) => Poll::Ready(Some(Err(err))), + Poll::Ready(None) => Poll::Ready(None), + Poll::Pending => Poll::Pending, + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + if self.batch_rx.is_closed() { + let len = self.batch_rx.len(); + (len, Some(len)) + } else { + (self.batch_rx.len(), None) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_common::{DataFusionError, assert_contains}; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryPool, UnboundedMemoryPool, + }; + use std::error::Error; + use std::fmt::Debug; + use std::time::Duration; + use tokio::time::timeout; + + #[tokio::test] + async fn buffers_only_some_messages() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); + let (_, res) = memory_pool_and_reservation(); + + let buffered = MemoryBufferedStream::new(input, 4, res); + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 2); + Ok(()) + } + + #[tokio::test] + async fn yields_all_messages() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); + let (_, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 4); + + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + finished(&mut buffered).await?; + Ok(()) + } + + #[tokio::test] + async fn yields_first_msg_even_if_big() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([25, 1, 2, 3]).map(Ok); + let (_, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 1); + pull_ok_msg(&mut buffered).await?; + Ok(()) + } + + #[tokio::test] + async fn memory_pool_kills_stream() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); + let (_, res) = bounded_memory_pool_and_reservation(7); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + let msg = pull_err_msg(&mut buffered).await?; + + assert_contains!(msg.to_string(), "Failed to allocate additional 4.0 B"); + Ok(()) + } + + #[tokio::test] + async fn memory_pool_does_not_kill_stream() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); + let (_, res) = bounded_memory_pool_and_reservation(7); + + let mut buffered = MemoryBufferedStream::new(input, 3, res); + wait_for_buffering().await; + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + finished(&mut buffered).await?; + Ok(()) + } + + #[tokio::test] + async fn messages_pass_even_if_all_exceed_limit() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([3, 3, 3, 3]).map(Ok); + let (_, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 2, res); + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 1); + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 1); + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 1); + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 1); + pull_ok_msg(&mut buffered).await?; + + wait_for_buffering().await; + finished(&mut buffered).await?; + Ok(()) + } + + #[tokio::test] + async fn errors_get_propagated() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([1, 2, 3, 4]).map(|v| { + if v == 3 { + return internal_err!("Error on 3"); + } + Ok(v) + }); + let (_, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + pull_err_msg(&mut buffered).await?; + + Ok(()) + } + + #[tokio::test] + async fn panic_in_input_is_propagated() -> Result<(), Box<dyn Error>> { + // A panic while polling the input must surface as a stream error, not a + // silent end-of-stream that drops the rest of the partition's output. + let input = futures::stream::iter([1, 2, 3, 4]).map(|v| { + if v == 3 { + panic!("boom on 3"); + } + Ok(v) + }); + let (_, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + let err = pull_err_msg(&mut buffered).await?; + assert_contains!(err.to_string(), "panicked"); + + Ok(()) + } + + #[tokio::test] + async fn memory_gets_released_if_stream_drops() -> Result<(), Box<dyn Error>> { + let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); + let (pool, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + assert_eq!(buffered.messages_queued(), 4); + assert_eq!(pool.reserved(), 10); + + pull_ok_msg(&mut buffered).await?; + assert_eq!(buffered.messages_queued(), 3); + assert_eq!(pool.reserved(), 9); + + pull_ok_msg(&mut buffered).await?; + assert_eq!(buffered.messages_queued(), 2); + assert_eq!(pool.reserved(), 7); + + drop(buffered); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + + fn memory_pool_and_reservation() -> (Arc<dyn MemoryPool>, MemoryReservation) { + let pool = Arc::new(UnboundedMemoryPool::default()) as _; + let reservation = MemoryConsumer::new("test").register(&pool); + (pool, reservation) + } + + fn bounded_memory_pool_and_reservation( + size: usize, + ) -> (Arc<dyn MemoryPool>, MemoryReservation) { + let pool = Arc::new(GreedyMemoryPool::new(size)) as _; + let reservation = MemoryConsumer::new("test").register(&pool); + (pool, reservation) + } + + async fn wait_for_buffering() { + // We do not have control over the spawned task, so the best we can do is to yield some + // cycles to the tokio runtime and let the task make progress on its own. + tokio::time::sleep(Duration::from_millis(1)).await; + } + + async fn pull_ok_msg<T: SizedMessage>( + buffered: &mut MemoryBufferedStream<T>, + ) -> Result<T, Box<dyn Error>> { + Ok(timeout(Duration::from_millis(1), buffered.next()) + .await? + .unwrap_or_else(|| internal_err!("Stream should not have finished"))?) + } + + async fn pull_err_msg<T: SizedMessage + Debug>( + buffered: &mut MemoryBufferedStream<T>, + ) -> Result<DataFusionError, Box<dyn Error>> { + Ok(timeout(Duration::from_millis(1), buffered.next()) + .await? + .map(|v| match v { + Ok(v) => internal_err!( + "Stream should not have failed, but succeeded with {v:?}" + ), + Err(err) => Ok(err), + }) + .unwrap_or_else(|| internal_err!("Stream should not have finished"))?) + } + + async fn finished<T: SizedMessage>( + buffered: &mut MemoryBufferedStream<T>, + ) -> Result<(), Box<dyn Error>> { + match timeout(Duration::from_millis(1), buffered.next()) + .await? + .is_none() + { + true => Ok(()), + false => internal_err!("Stream should have finished")?, + } + } + + impl SizedMessage for usize { + fn size(&self) -> usize { + *self + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/coalesce/mod.rs b/native/vendor/datafusion-physical-plan/src/coalesce/mod.rs new file mode 100644 index 0000000..ea1a87d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/coalesce/mod.rs
@@ -0,0 +1,375 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::RecordBatch; +use arrow::compute::BatchCoalescer; +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, assert_or_internal_err}; + +/// Concatenate multiple [`RecordBatch`]es and apply a limit +/// +/// See [`BatchCoalescer`] for more details on how this works. +#[derive(Debug)] +pub struct LimitedBatchCoalescer { + /// The arrow structure that builds the output batches + inner: BatchCoalescer, + /// Total number of rows returned so far + total_rows: usize, + /// Limit: maximum number of rows to fetch, `None` means fetch all rows + fetch: Option<usize>, + /// Indicates if the coalescer is finished + finished: bool, +} + +/// Status returned by [`LimitedBatchCoalescer::push_batch`] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PushBatchStatus { + /// The limit has **not** been reached, and more batches can be pushed + Continue, + /// The limit **has** been reached after processing this batch + /// The caller should call [`LimitedBatchCoalescer::finish`] + /// to flush any buffered rows and stop pushing more batches. + LimitReached, +} + +impl LimitedBatchCoalescer { + /// Create a new `BatchCoalescer` + /// + /// # Arguments + /// - `schema` - the schema of the output batches + /// - `target_batch_size` - the minimum number of rows for each + /// output batch (until limit reached) + /// - `fetch` - the maximum number of rows to fetch, `None` means fetch all rows + pub fn new( + schema: SchemaRef, + target_batch_size: usize, + fetch: Option<usize>, + ) -> Self { + Self { + inner: BatchCoalescer::new(schema, target_batch_size) + .with_biggest_coalesce_batch_size(Some(target_batch_size / 2)), + total_rows: 0, + fetch, + finished: false, + } + } + + /// Return the schema of the output batches + pub fn schema(&self) -> SchemaRef { + self.inner.schema() + } + + /// Pushes the next [`RecordBatch`] into the coalescer and returns its status. + /// + /// # Arguments + /// * `batch` - The [`RecordBatch`] to append. + /// + /// # Returns + /// * [`PushBatchStatus::Continue`] - More batches can still be pushed. + /// * [`PushBatchStatus::LimitReached`] - The row limit was reached after processing + /// this batch. The caller should call [`Self::finish`] before retrieving the + /// remaining buffered batches. + /// + /// # Errors + /// Returns an error if called after [`Self::finish`] or if the internal push + /// operation fails. + pub fn push_batch(&mut self, batch: RecordBatch) -> Result<PushBatchStatus> { + assert_or_internal_err!( + !self.finished, + "LimitedBatchCoalescer: cannot push batch after finish" + ); + + // if we are at the limit, return LimitReached + if let Some(fetch) = self.fetch { + // limit previously reached + if self.total_rows >= fetch { + return Ok(PushBatchStatus::LimitReached); + } + + // limit now reached + if self.total_rows + batch.num_rows() >= fetch { + // Limit is reached + let remaining_rows = fetch - self.total_rows; + debug_assert!(remaining_rows > 0); + + let batch_head = batch.slice(0, remaining_rows); + self.total_rows += batch_head.num_rows(); + self.inner.push_batch(batch_head)?; + return Ok(PushBatchStatus::LimitReached); + } + } + + // Limit not reached, push the entire batch + self.total_rows += batch.num_rows(); + self.inner.push_batch(batch)?; + + Ok(PushBatchStatus::Continue) + } + + /// Return true if there is no data buffered + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } + + /// Complete the current buffered batch and finish the coalescer + /// + /// Any subsequent calls to `push_batch()` will return an Err + pub fn finish(&mut self) -> Result<()> { + self.inner.finish_buffered_batch()?; + self.finished = true; + Ok(()) + } + + pub(crate) fn is_finished(&self) -> bool { + self.finished + } + + /// Return the next completed batch, if any + pub fn next_completed_batch(&mut self) -> Option<RecordBatch> { + self.inner.next_completed_batch() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::ops::Range; + use std::sync::Arc; + + use arrow::array::UInt32Array; + use arrow::compute::concat_batches; + use arrow::datatypes::{DataType, Field, Schema}; + + #[test] + fn test_coalesce() { + let batch = uint32_batch(0..8); + Test::new() + .with_batches(std::iter::repeat_n(batch, 10)) + // expected output is batches of exactly 21 rows (except for the final batch) + .with_target_batch_size(21) + .with_expected_output_sizes(vec![21, 21, 21, 17]) + .run() + } + + #[test] + fn test_coalesce_with_fetch_larger_than_input_size() { + let batch = uint32_batch(0..8); + Test::new() + .with_batches(std::iter::repeat_n(batch, 10)) + // input is 10 batches x 8 rows (80 rows) with fetch limit of 100 + // expected to behave the same as `test_concat_batches` + .with_target_batch_size(21) + .with_fetch(Some(100)) + .with_expected_output_sizes(vec![21, 21, 21, 17]) + .run(); + } + + #[test] + fn test_coalesce_with_fetch_less_than_input_size() { + let batch = uint32_batch(0..8); + Test::new() + .with_batches(std::iter::repeat_n(batch, 10)) + // input is 10 batches x 8 rows (80 rows) with fetch limit of 50 + .with_target_batch_size(21) + .with_fetch(Some(50)) + .with_expected_output_sizes(vec![21, 21, 8]) + .run(); + } + + #[test] + fn test_coalesce_with_fetch_less_than_target_and_no_remaining_rows() { + let batch = uint32_batch(0..8); + Test::new() + .with_batches(std::iter::repeat_n(batch, 10)) + // input is 10 batches x 8 rows (80 rows) with fetch limit of 48 + .with_target_batch_size(24) + .with_fetch(Some(48)) + .with_expected_output_sizes(vec![24, 24]) + .run(); + } + + #[test] + fn test_coalesce_with_fetch_less_target_batch_size() { + let batch = uint32_batch(0..8); + Test::new() + .with_batches(std::iter::repeat_n(batch, 10)) + // input is 10 batches x 8 rows (80 rows) with fetch limit of 10 + .with_target_batch_size(21) + .with_fetch(Some(10)) + .with_expected_output_sizes(vec![10]) + .run(); + } + + #[test] + fn test_coalesce_single_large_batch_over_fetch() { + let large_batch = uint32_batch(0..100); + Test::new() + .with_batch(large_batch) + .with_target_batch_size(20) + .with_fetch(Some(7)) + .with_expected_output_sizes(vec![7]) + .run() + } + + /// Test for [`LimitedBatchCoalescer`] + /// + /// Pushes the input batches to the coalescer and verifies that the resulting + /// batches have the expected number of rows and contents. + #[derive(Debug, Clone, Default)] + struct Test { + /// Batches to feed to the coalescer. Tests must have at least one + /// schema + input_batches: Vec<RecordBatch>, + /// Expected output sizes of the resulting batches + expected_output_sizes: Vec<usize>, + /// target batch size + target_batch_size: usize, + /// Fetch (limit) + fetch: Option<usize>, + } + + impl Test { + fn new() -> Self { + Self::default() + } + + /// Set the target batch size + fn with_target_batch_size(mut self, target_batch_size: usize) -> Self { + self.target_batch_size = target_batch_size; + self + } + + /// Set the fetch (limit) + fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + /// Extend the input batches with `batch` + fn with_batch(mut self, batch: RecordBatch) -> Self { + self.input_batches.push(batch); + self + } + + /// Extends the input batches with `batches` + fn with_batches( + mut self, + batches: impl IntoIterator<Item = RecordBatch>, + ) -> Self { + self.input_batches.extend(batches); + self + } + + /// Extends `sizes` to expected output sizes + fn with_expected_output_sizes( + mut self, + sizes: impl IntoIterator<Item = usize>, + ) -> Self { + self.expected_output_sizes.extend(sizes); + self + } + + /// Runs the test -- see documentation on [`Test`] for details + fn run(self) { + let Self { + input_batches, + target_batch_size, + fetch, + expected_output_sizes, + } = self; + + let schema = input_batches[0].schema(); + + // create a single large input batch for output comparison + let single_input_batch = concat_batches(&schema, &input_batches).unwrap(); + + let mut coalescer = + LimitedBatchCoalescer::new(Arc::clone(&schema), target_batch_size, fetch); + + let mut output_batches = vec![]; + for batch in input_batches { + match coalescer.push_batch(batch).unwrap() { + PushBatchStatus::Continue => { + // continue pushing batches + } + PushBatchStatus::LimitReached => { + break; + } + } + } + coalescer.finish().unwrap(); + while let Some(batch) = coalescer.next_completed_batch() { + output_batches.push(batch); + } + + let actual_output_sizes: Vec<usize> = + output_batches.iter().map(|b| b.num_rows()).collect(); + assert_eq!( + expected_output_sizes, actual_output_sizes, + "Unexpected number of rows in output batches\n\ + Expected\n{expected_output_sizes:#?}\nActual:{actual_output_sizes:#?}" + ); + + // make sure we got the expected number of output batches and content + let mut starting_idx = 0; + assert_eq!(expected_output_sizes.len(), output_batches.len()); + for (i, (expected_size, batch)) in + expected_output_sizes.iter().zip(output_batches).enumerate() + { + assert_eq!( + *expected_size, + batch.num_rows(), + "Unexpected number of rows in Batch {i}" + ); + + // compare the contents of the batch (using `==` compares the + // underlying memory layout too) + let expected_batch = + single_input_batch.slice(starting_idx, *expected_size); + let batch_strings = batch_to_pretty_strings(&batch); + let expected_batch_strings = batch_to_pretty_strings(&expected_batch); + let batch_strings = batch_strings.lines().collect::<Vec<_>>(); + let expected_batch_strings = + expected_batch_strings.lines().collect::<Vec<_>>(); + assert_eq!( + expected_batch_strings, batch_strings, + "Unexpected content in Batch {i}:\ + \n\nExpected:\n{expected_batch_strings:#?}\n\nActual:\n{batch_strings:#?}" + ); + starting_idx += *expected_size; + } + } + } + + /// Return a batch of UInt32 with the specified range + fn uint32_batch(range: Range<u32>) -> RecordBatch { + let schema = + Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])); + + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from_iter_values(range))], + ) + .unwrap() + } + + fn batch_to_pretty_strings(batch: &RecordBatch) -> String { + arrow::util::pretty::pretty_format_batches(std::slice::from_ref(batch)) + .unwrap() + .to_string() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/coalesce_batches.rs b/native/vendor/datafusion-physical-plan/src/coalesce_batches.rs new file mode 100644 index 0000000..cb0f9b2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/coalesce_batches.rs
@@ -0,0 +1,459 @@ +// 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. + +//! [`CoalesceBatchesExec`] combines small batches into larger batches. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics}; +use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::EmptyRecordBatchStream; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, RecordBatchStream, + ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, +}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; + +use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; +use crate::filter_pushdown::{ + ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::sort_pushdown::SortOrderPushdownResult; +use datafusion_common::config::ConfigOptions; +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use futures::ready; +use futures::stream::{Stream, StreamExt}; + +/// `CoalesceBatchesExec` combines small batches into larger batches for more +/// efficient vectorized processing by later operators. +/// +/// The operator buffers batches until it collects `target_batch_size` rows and +/// then emits a single concatenated batch. When only a limited number of rows +/// are necessary (specified by the `fetch` parameter), the operator will stop +/// buffering and returns the final batch once the number of collected rows +/// reaches the `fetch` value. +/// +/// See [`LimitedBatchCoalescer`] for more information +#[deprecated( + since = "52.0.0", + note = "We now use BatchCoalescer from arrow-rs instead of a dedicated operator" +)] +#[derive(Debug, Clone)] +pub struct CoalesceBatchesExec { + /// The input plan + input: Arc<dyn ExecutionPlan>, + /// Minimum number of rows for coalescing batches + target_batch_size: usize, + /// Maximum number of rows to fetch, `None` means fetching all rows + fetch: Option<usize>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + cache: Arc<PlanProperties>, +} + +#[expect(deprecated)] +impl CoalesceBatchesExec { + /// Create a new CoalesceBatchesExec + pub fn new(input: Arc<dyn ExecutionPlan>, target_batch_size: usize) -> Self { + let cache = Self::compute_properties(&input); + Self { + input, + target_batch_size, + fetch: None, + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + } + } + + /// Update fetch with the argument + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + /// The input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Minimum number of rows for coalesces batches + pub fn target_batch_size(&self) -> usize { + self.target_batch_size + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties { + // The coalesce batches operator does not make any changes to the + // partitioning of its input. + PlanProperties::new( + input.equivalence_properties().clone(), // Equivalence Properties + input.output_partitioning().clone(), // Output Partitioning + input.pipeline_behavior(), + input.boundedness(), + ) + } +} + +#[expect(deprecated)] +impl DisplayAs for CoalesceBatchesExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "CoalesceBatchesExec: target_batch_size={}", + self.target_batch_size, + )?; + if let Some(fetch) = self.fetch { + write!(f, ", fetch={fetch}")?; + }; + + Ok(()) + } + DisplayFormatType::TreeRender => { + writeln!(f, "target_batch_size={}", self.target_batch_size)?; + if let Some(fetch) = self.fetch { + write!(f, "limit={fetch}")?; + }; + Ok(()) + } + } + } +} + +#[expect(deprecated)] +impl ExecutionPlan for CoalesceBatchesExec { + fn name(&self) -> &'static str { + "CoalesceBatchesExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(CoalesceBatchesStream { + input: self.input.execute(partition, context)?, + coalescer: LimitedBatchCoalescer::new( + self.input.schema(), + self.target_batch_size, + self.fetch, + ), + baseline_metrics: BaselineMetrics::new(&self.metrics, partition), + completed: false, + })) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + Some(Arc::new(CoalesceBatchesExec { + input: Arc::clone(&self.input), + target_batch_size: self.target_batch_size, + fetch: limit, + metrics: self.metrics.clone(), + cache: Arc::clone(&self.cache), + })) + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + match self.input.try_swapping_with_projection(projection)? { + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), + None => Ok(None), + } + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + FilterDescription::from_children(parent_filters, &self.children()) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + Ok(FilterPushdownPropagation::if_all(child_pushdown_result)) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + // CoalesceBatchesExec is transparent for sort ordering - it preserves order + // Delegate to the child and wrap with a new CoalesceBatchesExec + self.input.try_pushdown_sort(order)?.try_map(|new_input| { + Ok(Arc::new( + CoalesceBatchesExec::new(new_input, self.target_batch_size) + .with_fetch(self.fetch), + ) as Arc<dyn ExecutionPlan>) + }) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches( + Box::new(protobuf::CoalesceBatchesExecNode { + input: Some(Box::new(input)), + target_batch_size: self.target_batch_size() as u32, + fetch: self.fetch().map(|n| n as u32), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +#[expect(deprecated)] +impl CoalesceBatchesExec { + /// Reconstruct a [`CoalesceBatchesExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. The child plan is decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let coalesce_batches = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches, + "CoalesceBatchesExec", + ); + let input = ctx.decode_required_child( + coalesce_batches.input.as_deref(), + "CoalesceBatchesExec", + "input", + )?; + Ok(Arc::new( + CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) + .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), + )) + } +} + +/// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details. +struct CoalesceBatchesStream { + /// The input plan + input: SendableRecordBatchStream, + /// Buffer for combining batches + coalescer: LimitedBatchCoalescer, + /// Execution metrics + baseline_metrics: BaselineMetrics, + /// is the input stream exhausted or limit reached? + completed: bool, +} + +impl Stream for CoalesceBatchesStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let poll = self.poll_next_inner(cx); + self.baseline_metrics.record_poll(poll) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + // we can't predict the size of incoming batches so re-use the size hint from the input + self.input.size_hint() + } +} + +impl CoalesceBatchesStream { + fn poll_next_inner( + self: &mut Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + let cloned_time = self.baseline_metrics.elapsed_compute().clone(); + loop { + // If there is any completed batch ready, return it + if let Some(batch) = self.coalescer.next_completed_batch() { + return Poll::Ready(Some(Ok(batch))); + } + if self.completed { + // If input is done and no batches are ready, return None to signal end of stream. + return Poll::Ready(None); + } + // Attempt to pull the next batch from the input stream. + let input_batch = ready!(self.input.poll_next_unpin(cx)); + // Start timing the operation. The timer records time upon being dropped. + let _timer = cloned_time.timer(); + + match input_batch { + None => { + // Input stream is exhausted, finalize any remaining batches + self.completed = true; + self.input = + Box::pin(EmptyRecordBatchStream::new(self.coalescer.schema())); + self.coalescer.finish()?; + } + Some(Ok(batch)) => { + match self.coalescer.push_batch(batch)? { + PushBatchStatus::Continue => { + // Keep pushing more batches + } + PushBatchStatus::LimitReached => { + // limit was reached, so stop early + self.completed = true; + self.input = Box::pin(EmptyRecordBatchStream::new( + self.coalescer.schema(), + )); + self.coalescer.finish()?; + } + } + } + // Error case + other => return Poll::Ready(other), + } + } + } +} + +impl RecordBatchStream for CoalesceBatchesStream { + fn schema(&self) -> SchemaRef { + self.coalescer.schema() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/coalesce_partitions.rs b/native/vendor/datafusion-physical-plan/src/coalesce_partitions.rs new file mode 100644 index 0000000..6f58eb2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/coalesce_partitions.rs
@@ -0,0 +1,657 @@ +// 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. + +//! Defines the merge plan for executing partitions in parallel and then merging the results +//! into a single partition + +use std::sync::Arc; + +use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use super::stream::{ObservedStream, RecordBatchReceiverStream}; +use super::{ + DisplayAs, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, + Statistics, +}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; +use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; +use crate::projection::{ProjectionExec, make_with_child}; +use crate::sort_pushdown::SortOrderPushdownResult; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, validate_child_count, +}; +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; + +/// Merge execution plan executes partitions in parallel and combines them into a single +/// partition. No guarantees are made about the order of the resulting partition. +#[derive(Debug, Clone)] +pub struct CoalescePartitionsExec { + /// Input execution plan + input: Arc<dyn ExecutionPlan>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + cache: Arc<PlanProperties>, + /// Optional number of rows to fetch. Stops producing rows after this fetch + pub(crate) fetch: Option<usize>, +} + +impl CoalescePartitionsExec { + /// Create a new CoalescePartitionsExec + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + let cache = Self::compute_properties(&input); + CoalescePartitionsExec { + input, + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + fetch: None, + } + } + + /// Update fetch with the argument + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + /// Input execution plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties { + let input_partitions = input.output_partitioning().partition_count(); + let (drive, scheduling) = if input_partitions > 1 { + (EvaluationType::Eager, SchedulingType::Cooperative) + } else { + ( + input.properties().evaluation_type, + input.properties().scheduling_type, + ) + }; + + // Coalescing partitions loses existing orderings: + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.clear_orderings(); + eq_properties.clear_per_partition_constants(); + PlanProperties::new( + eq_properties, // Equivalence Properties + Partitioning::UnknownPartitioning(1), // Output Partitioning + input.pipeline_behavior(), + input.boundedness(), + ) + .with_evaluation_type(drive) + .with_scheduling_type(scheduling) + } +} + +impl DisplayAs for CoalescePartitionsExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => match self.fetch { + Some(fetch) => { + write!(f, "CoalescePartitionsExec: fetch={fetch}") + } + None => write!(f, "CoalescePartitionsExec"), + }, + DisplayFormatType::TreeRender => match self.fetch { + Some(fetch) => { + write!(f, "limit: {fetch}") + } + None => write!(f, ""), + }, + } + } +} + +impl ExecutionPlan for CoalescePartitionsExec { + fn name(&self) -> &'static str { + "CoalescePartitionsExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); + plan.fetch = self.fetch; + Ok(Arc::new(plan)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + // CoalescePartitionsExec produces a single partition + assert_eq_or_internal_err!( + partition, + 0, + "CoalescePartitionsExec invalid partition {partition}" + ); + + let input_partitions = self.input.output_partitioning().partition_count(); + match input_partitions { + 0 => internal_err!( + "CoalescePartitionsExec requires at least one input partition" + ), + 1 => { + // single-partition path: execute child directly, but ensure fetch is respected + // (wrap with ObservedStream only if fetch is present so we don't add overhead otherwise) + let child_stream = self.input.execute(0, context)?; + if self.fetch.is_some() { + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + return Ok(Box::pin(ObservedStream::new( + child_stream, + baseline_metrics, + self.fetch, + ))); + } + Ok(child_stream) + } + _ => { + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + // record the (very) minimal work done so that + // elapsed_compute is not reported as 0 + let elapsed_compute = baseline_metrics.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); + + // use a stream that allows each sender to put in at + // least one result in an attempt to maximize + // parallelism. + let mut builder = + RecordBatchReceiverStream::builder(self.schema(), input_partitions); + + // spawn independent tasks whose resulting streams (of batches) + // are sent to the channel for consumption. + for part_i in 0..input_partitions { + builder.run_input( + Arc::clone(&self.input), + part_i, + Arc::clone(&context), + ); + } + + let stream = builder.build(); + Ok(Box::pin(ObservedStream::new( + stream, + baseline_metrics, + self.fetch, + ))) + } + } + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(None)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + /// Tries to swap `projection` with its input, which is known to be a + /// [`CoalescePartitionsExec`]. If possible, performs the swap and returns + /// [`CoalescePartitionsExec`] as the top plan. Otherwise, returns `None`. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection does not narrow the schema, we should not try to push it down: + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + // CoalescePartitionsExec always has a single child, so zero indexing is safe. + make_with_child(projection, projection.input().children()[0]).map(|e| { + if self.fetch.is_some() { + let mut plan = CoalescePartitionsExec::new(e); + plan.fetch = self.fetch; + Some(Arc::new(plan) as _) + } else { + Some(Arc::new(CoalescePartitionsExec::new(e)) as _) + } + }) + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + Some(Arc::new(CoalescePartitionsExec { + input: Arc::clone(&self.input), + fetch: limit, + metrics: self.metrics.clone(), + cache: Arc::clone(&self.cache), + })) + } + + fn with_preserve_order( + &self, + preserve_order: bool, + ) -> Option<Arc<dyn ExecutionPlan>> { + self.input + .with_preserve_order(preserve_order) + .and_then(|new_input| { + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) + .ok() + }) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + FilterDescription::from_children(parent_filters, &self.children()) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + // CoalescePartitionsExec merges multiple partitions into one, which loses + // global ordering. However, we can still push the sort requirement down + // to optimize individual partitions - the Sort operator above will handle + // the global ordering. + // + // Note: The result will always be at most Inexact (never Exact) when there + // are multiple partitions, because merging destroys global ordering. + let result = self.input.try_pushdown_sort(order)?; + + // If we have multiple partitions, we can't return Exact even if the + // underlying source claims Exact - merging destroys global ordering + let has_multiple_partitions = + self.input.output_partitioning().partition_count() > 1; + + result + .try_map(|new_input| { + Ok( + Arc::new( + CoalescePartitionsExec::new(new_input).with_fetch(self.fetch), + ) as Arc<dyn ExecutionPlan>, + ) + }) + .map(|r| { + if has_multiple_partitions { + // Downgrade Exact to Inexact when merging multiple partitions + r.into_inexact() + } else { + r + } + }) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Merge(Box::new( + protobuf::CoalescePartitionsExecNode { + input: Some(Box::new(input)), + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CoalescePartitionsExec { + /// Reconstruct a [`CoalescePartitionsExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. Note the protobuf + /// variant is named `Merge` (node [`CoalescePartitionsExecNode`]). + /// + /// [`CoalescePartitionsExecNode`]: datafusion_proto_models::protobuf::CoalescePartitionsExecNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let merge = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Merge, + "CoalescePartitionsExec", + ); + let input = ctx.decode_required_child( + merge.input.as_deref(), + "CoalescePartitionsExec", + "input", + )?; + Ok(Arc::new( + CoalescePartitionsExec::new(input) + .with_fetch(merge.fetch.map(|f| f as usize)), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test::exec::{ + BarrierExec, BlockingExec, PanicExec, assert_strong_count_converges_to_zero, + }; + use crate::test::{self, assert_is_pending}; + use crate::{collect, common}; + + use std::time::Duration; + + use arrow::array::RecordBatch; + use arrow::datatypes::{DataType, Field, Schema}; + + use futures::FutureExt; + + #[tokio::test] + async fn merge() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + let num_partitions = 4; + let csv = test::scan_partitioned(num_partitions); + + // input should have 4 partitions + assert_eq!(csv.output_partitioning().partition_count(), num_partitions); + + let merge = CoalescePartitionsExec::new(csv); + + // output of CoalescePartitionsExec should have a single partition + assert_eq!( + merge.properties().output_partitioning().partition_count(), + 1 + ); + + // the result should contain 4 batches (one per input partition) + let iter = merge.execute(0, task_ctx)?; + let batches = common::collect(iter).await?; + assert_eq!(batches.len(), num_partitions); + + // there should be a total of 400 rows (100 per each partition) + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(row_count, 400); + + Ok(()) + } + + #[tokio::test] + async fn drops_input_plan_after_input_streams_start() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + let input_partitions = 2; + let batch = RecordBatch::new_empty(Arc::clone(&schema)); + let input = Arc::new( + BarrierExec::new(vec![vec![batch]; input_partitions], schema) + .without_start_barrier() + .with_finish_barrier() + .with_log(false), + ); + let refs = Arc::downgrade(&input); + + let input_plan: Arc<BarrierExec> = Arc::clone(&input); + let coalesce = CoalescePartitionsExec::new(input_plan); + let stream = coalesce.execute(0, task_ctx)?; + drop(coalesce); + + tokio::time::timeout(Duration::from_secs(5), async { + // Why not `wait_finish` here: that releases the barrier which lets the input tasks + // finish, which drops the input Arcs and hides the bug. + while !input.is_finish_barrier_reached() { + tokio::task::yield_now().await; + } + }) + .await + .expect("input streams should reach pending"); + + drop(input); + + assert_strong_count_converges_to_zero(refs).await; + + drop(stream); + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 2)); + let refs = blocking_exec.refs(); + let coalesce_partitions_exec = + Arc::new(CoalescePartitionsExec::new(blocking_exec)); + + let fut = collect(coalesce_partitions_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + #[should_panic(expected = "PanickingStream did panic")] + async fn test_panic() { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let panicking_exec = Arc::new(PanicExec::new(Arc::clone(&schema), 2)); + let coalesce_partitions_exec = + Arc::new(CoalescePartitionsExec::new(panicking_exec)); + + collect(coalesce_partitions_exec, task_ctx).await.unwrap(); + } + + #[tokio::test] + async fn test_single_partition_with_fetch() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + // Use existing scan_partitioned with 1 partition (returns 100 rows per partition) + let input = test::scan_partitioned(1); + + // Test with fetch=3 + let coalesce = CoalescePartitionsExec::new(input).with_fetch(Some(3)); + + let stream = coalesce.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(row_count, 3, "Should only return 3 rows due to fetch=3"); + + Ok(()) + } + + #[tokio::test] + async fn test_multi_partition_with_fetch_one() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + // Create 4 partitions, each with 100 rows + // This simulates the real-world scenario where each partition has data + let input = test::scan_partitioned(4); + + // Test with fetch=1 (the original bug: was returning multiple rows instead of 1) + let coalesce = CoalescePartitionsExec::new(input).with_fetch(Some(1)); + + let stream = coalesce.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!( + row_count, 1, + "Should only return 1 row due to fetch=1, not one per partition" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_single_partition_without_fetch() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + // Use scan_partitioned with 1 partition + let input = test::scan_partitioned(1); + + // Test without fetch (should return all rows) + let coalesce = CoalescePartitionsExec::new(input); + + let stream = coalesce.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!( + row_count, 100, + "Should return all 100 rows when fetch is None" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_single_partition_fetch_larger_than_batch() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + // Use scan_partitioned with 1 partition (returns 100 rows) + let input = test::scan_partitioned(1); + + // Test with fetch larger than available rows + let coalesce = CoalescePartitionsExec::new(input).with_fetch(Some(200)); + + let stream = coalesce.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!( + row_count, 100, + "Should return all available rows (100) when fetch (200) is larger" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_multi_partition_fetch_exact_match() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + // Create 4 partitions, each with 100 rows + let num_partitions = 4; + let csv = test::scan_partitioned(num_partitions); + + // Test with fetch=400 (exactly all rows) + let coalesce = CoalescePartitionsExec::new(csv).with_fetch(Some(400)); + + let stream = coalesce.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(row_count, 400, "Should return exactly 400 rows"); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/column_rewriter.rs b/native/vendor/datafusion-physical-plan/src/column_rewriter.rs new file mode 100644 index 0000000..2df95cd --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/column_rewriter.rs
@@ -0,0 +1,382 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use datafusion_common::{ + DataFusionError, HashMap, + tree_node::{Transformed, TreeNodeRecursion, TreeNodeRewriter}, +}; +use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; + +/// Rewrite column references in a physical expr according to a mapping. +/// +/// This rewriter traverses the expression tree and replaces [`Column`] nodes +/// with the corresponding expression found in the `column_map`. +/// +/// If a column is found in the map, it is replaced by the mapped expression. +/// If a column is NOT found in the map, a `DataFusionError::Internal` is +/// returned. +pub struct PhysicalColumnRewriter<'a> { + /// Mapping from original column to new column. + pub column_map: &'a HashMap<Column, Arc<dyn PhysicalExpr>>, +} + +impl<'a> PhysicalColumnRewriter<'a> { + /// Create a new PhysicalColumnRewriter with the given column mapping. + pub fn new(column_map: &'a HashMap<Column, Arc<dyn PhysicalExpr>>) -> Self { + Self { column_map } + } +} + +impl<'a> TreeNodeRewriter for PhysicalColumnRewriter<'a> { + type Node = Arc<dyn PhysicalExpr>; + + fn f_down( + &mut self, + node: Self::Node, + ) -> datafusion_common::Result<Transformed<Self::Node>> { + if let Some(column) = node.downcast_ref::<Column>() { + if let Some(new_column) = self.column_map.get(column) { + // jump to prevent rewriting the new sub-expression again + return Ok(Transformed::new( + Arc::clone(new_column), + true, + TreeNodeRecursion::Jump, + )); + } else { + // Column not found in mapping + return Err(DataFusionError::Internal(format!( + "Column {column:?} not found in column mapping {:?}", + self.column_map + ))); + } + } + Ok(Transformed::no(node)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{Result, tree_node::TreeNode}; + use datafusion_physical_expr::{ + PhysicalExpr, + expressions::{Column, binary, col, lit}, + }; + + /// Helper function to create a test schema + fn create_test_schema() -> Arc<Schema> { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + Field::new("new_col", DataType::Int32, true), + Field::new("inner_col", DataType::Int32, true), + Field::new("another_col", DataType::Int32, true), + ])) + } + + /// Helper function to create a complex nested expression with multiple columns + /// Create: (col_a + col_b) * (col_c - col_d) + col_e + fn create_complex_expression(schema: &Schema) -> Arc<dyn PhysicalExpr> { + let col_a = col("a", schema).unwrap(); + let col_b = col("b", schema).unwrap(); + let col_c = col("c", schema).unwrap(); + let col_d = col("d", schema).unwrap(); + let col_e = col("e", schema).unwrap(); + + let add_expr = + binary(col_a, datafusion_expr::Operator::Plus, col_b, schema).unwrap(); + let sub_expr = + binary(col_c, datafusion_expr::Operator::Minus, col_d, schema).unwrap(); + let mul_expr = binary( + add_expr, + datafusion_expr::Operator::Multiply, + sub_expr, + schema, + ) + .unwrap(); + binary(mul_expr, datafusion_expr::Operator::Plus, col_e, schema).unwrap() + } + + /// Helper function to create a deeply nested expression + /// Create: col_a + (col_b + (col_c + (col_d + col_e))) + fn create_deeply_nested_expression(schema: &Schema) -> Arc<dyn PhysicalExpr> { + let col_a = col("a", schema).unwrap(); + let col_b = col("b", schema).unwrap(); + let col_c = col("c", schema).unwrap(); + let col_d = col("d", schema).unwrap(); + let col_e = col("e", schema).unwrap(); + + let inner1 = + binary(col_d, datafusion_expr::Operator::Plus, col_e, schema).unwrap(); + let inner2 = + binary(col_c, datafusion_expr::Operator::Plus, inner1, schema).unwrap(); + let inner3 = + binary(col_b, datafusion_expr::Operator::Plus, inner2, schema).unwrap(); + binary(col_a, datafusion_expr::Operator::Plus, inner3, schema).unwrap() + } + + #[test] + fn test_simple_column_replacement_with_jump() -> Result<()> { + let schema = create_test_schema(); + + // Test that Jump prevents re-processing of replaced columns + let mut column_map = HashMap::new(); + column_map.insert(Column::new_with_schema("a", &schema).unwrap(), lit(42i32)); + column_map.insert( + Column::new_with_schema("b", &schema).unwrap(), + lit("replaced_b"), + ); + column_map.insert( + Column::new_with_schema("c", &schema).unwrap(), + col("c", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("d", &schema).unwrap(), + col("d", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("e", &schema).unwrap(), + col("e", &schema).unwrap(), + ); + + let mut rewriter = PhysicalColumnRewriter::new(&column_map); + let expr = create_complex_expression(&schema); + + let result = expr.rewrite(&mut rewriter)?; + + // Verify the transformation occurred + assert!(result.transformed); + + assert_eq!( + format!("{}", result.data), + "(42 + replaced_b) * (c@2 - d@3) + e@4" + ); + + Ok(()) + } + + #[test] + fn test_nested_column_replacement_with_jump() -> Result<()> { + let schema = create_test_schema(); + // Test Jump behavior with deeply nested expressions + let mut column_map = HashMap::new(); + // Replace col_c with a complex expression containing new columns + let replacement_expr = binary( + lit(100i32), + datafusion_expr::Operator::Plus, + col("new_col", &schema).unwrap(), + &schema, + ) + .unwrap(); + column_map.insert( + Column::new_with_schema("c", &schema).unwrap(), + replacement_expr, + ); + column_map.insert( + Column::new_with_schema("a", &schema).unwrap(), + col("a", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("b", &schema).unwrap(), + col("b", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("d", &schema).unwrap(), + col("d", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("e", &schema).unwrap(), + col("e", &schema).unwrap(), + ); + + let mut rewriter = PhysicalColumnRewriter::new(&column_map); + let expr = create_deeply_nested_expression(&schema); + + let result = expr.rewrite(&mut rewriter)?; + + // Verify transformation occurred + assert!(result.transformed); + + assert_eq!( + format!("{}", result.data), + "a@0 + b@1 + 100 + new_col@5 + d@3 + e@4" + ); + + Ok(()) + } + + #[test] + fn test_circular_reference_prevention() -> Result<()> { + let schema = create_test_schema(); + // Test that Jump prevents infinite recursion with circular references + let mut column_map = HashMap::new(); + + // Create a circular reference: col_a -> col_b -> col_a (but Jump should prevent the second visit) + column_map.insert( + Column::new_with_schema("a", &schema).unwrap(), + col("b", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("b", &schema).unwrap(), + col("a", &schema).unwrap(), + ); + + let mut rewriter = PhysicalColumnRewriter::new(&column_map); + + // Start with an expression containing col_a + let expr = binary( + col("a", &schema).unwrap(), + datafusion_expr::Operator::Plus, + col("b", &schema).unwrap(), + &schema, + ) + .unwrap(); + + let result = expr.rewrite(&mut rewriter)?; + + // Verify transformation occurred + assert!(result.transformed); + + assert_eq!(format!("{}", result.data), "b@1 + a@0"); + + Ok(()) + } + + #[test] + fn test_multiple_replacements_in_same_expression() -> Result<()> { + let schema = create_test_schema(); + // Test multiple column replacements in the same complex expression + let mut column_map = HashMap::new(); + + // Replace multiple columns with literals + column_map.insert(Column::new_with_schema("a", &schema).unwrap(), lit(10i32)); + column_map.insert(Column::new_with_schema("c", &schema).unwrap(), lit(20i32)); + column_map.insert(Column::new_with_schema("e", &schema).unwrap(), lit(30i32)); + column_map.insert( + Column::new_with_schema("b", &schema).unwrap(), + col("b", &schema).unwrap(), + ); + column_map.insert( + Column::new_with_schema("d", &schema).unwrap(), + col("d", &schema).unwrap(), + ); + + let mut rewriter = PhysicalColumnRewriter::new(&column_map); + let expr = create_complex_expression(&schema); // (col_a + col_b) * (col_c - col_d) + col_e + + let result = expr.rewrite(&mut rewriter)?; + + // Verify transformation occurred + assert!(result.transformed); + assert_eq!(format!("{}", result.data), "(10 + b@1) * (20 - d@3) + 30"); + + Ok(()) + } + + #[test] + fn test_jump_with_complex_replacement_expression() -> Result<()> { + let schema = create_test_schema(); + // Test Jump behavior when replacing with very complex expressions + let mut column_map = HashMap::new(); + + // Replace col_a with a complex nested expression + let inner_expr = binary( + lit(5i32), + datafusion_expr::Operator::Multiply, + col("a", &schema).unwrap(), + &schema, + ) + .unwrap(); + let middle_expr = binary( + inner_expr, + datafusion_expr::Operator::Plus, + lit(3i32), + &schema, + ) + .unwrap(); + let complex_replacement = binary( + middle_expr, + datafusion_expr::Operator::Minus, + col("another_col", &schema).unwrap(), + &schema, + ) + .unwrap(); + + column_map.insert( + Column::new_with_schema("a", &schema).unwrap(), + complex_replacement, + ); + column_map.insert( + Column::new_with_schema("b", &schema).unwrap(), + col("b", &schema).unwrap(), + ); + + let mut rewriter = PhysicalColumnRewriter::new(&column_map); + + // Create expression: col_a + col_b + let expr = binary( + col("a", &schema).unwrap(), + datafusion_expr::Operator::Plus, + col("b", &schema).unwrap(), + &schema, + ) + .unwrap(); + + let result = expr.rewrite(&mut rewriter)?; + + assert_eq!( + format!("{}", result.data), + "5 * a@0 + 3 - another_col@7 + b@1" + ); + + // Verify transformation occurred + assert!(result.transformed); + + Ok(()) + } + + #[test] + fn test_unmapped_columns_detection() -> Result<()> { + let schema = create_test_schema(); + let mut column_map = HashMap::new(); + + // Only map col_a, leave col_b unmapped + column_map.insert(Column::new_with_schema("a", &schema).unwrap(), lit(42i32)); + + let mut rewriter = PhysicalColumnRewriter::new(&column_map); + + // Create expression: col_a + col_b + let expr = binary( + col("a", &schema).unwrap(), + datafusion_expr::Operator::Plus, + col("b", &schema).unwrap(), + &schema, + ) + .unwrap(); + + let err = expr.rewrite(&mut rewriter).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(_))); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/common.rs b/native/vendor/datafusion-physical-plan/src/common.rs new file mode 100644 index 0000000..734ec96 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/common.rs
@@ -0,0 +1,623 @@ +// 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. + +//! Defines common code used in execution plans + +use std::fs; +use std::fs::metadata; +use std::sync::Arc; + +use super::SendableRecordBatchStream; +use crate::expressions::{CastExpr, Column}; +use crate::projection::{ProjectionExec, ProjectionExpr}; +use crate::stream::RecordBatchReceiverStream; +use crate::{ColumnStatistics, ExecutionPlan, Statistics}; + +use arrow::array::Array; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::stats::Precision; +use datafusion_common::{Result, plan_err}; +use datafusion_execution::memory_pool::MemoryReservation; + +use futures::{StreamExt, TryStreamExt}; + +/// [`MemoryReservation`] used across query execution streams +pub(crate) type SharedMemoryReservation = Arc<MemoryReservation>; + +/// Create a vector of record batches from a stream +pub async fn collect(stream: SendableRecordBatchStream) -> Result<Vec<RecordBatch>> { + stream.try_collect::<Vec<_>>().await +} + +/// Recursively builds a list of files in a directory with a given extension +pub fn build_checked_file_list(dir: &str, ext: &str) -> Result<Vec<String>> { + let mut filenames: Vec<String> = Vec::new(); + build_file_list_recurse(dir, &mut filenames, ext)?; + if filenames.is_empty() { + return plan_err!("No files found at {dir} with file extension {ext}"); + } + Ok(filenames) +} + +/// Recursively builds a list of files in a directory with a given extension +pub fn build_file_list(dir: &str, ext: &str) -> Result<Vec<String>> { + let mut filenames: Vec<String> = Vec::new(); + build_file_list_recurse(dir, &mut filenames, ext)?; + Ok(filenames) +} + +/// Recursively build a list of files in a directory with a given extension with an accumulator list +fn build_file_list_recurse( + dir: &str, + filenames: &mut Vec<String>, + ext: &str, +) -> Result<()> { + let metadata = metadata(dir)?; + if metadata.is_file() { + if dir.ends_with(ext) { + filenames.push(dir.to_string()); + } + } else { + for entry in fs::read_dir(dir)? { + let entry = entry?; + let path = entry.path(); + if let Some(path_name) = path.to_str() { + if path.is_dir() { + build_file_list_recurse(path_name, filenames, ext)?; + } else if path_name.ends_with(ext) { + filenames.push(path_name.to_string()); + } + } else { + return plan_err!("Invalid path"); + } + } + } + Ok(()) +} + +/// Align `input`'s physical plan schema with `expected_schema`. +/// +/// This helper is intended for operators that combine independently planned children but +/// expose a single declared output schema. It returns `input` unchanged when schemas already +/// match exactly. Otherwise, it validates that projection can safely produce the expected +/// schema, then wraps `input` in a [`ProjectionExec`] that keeps columns in their existing +/// positional order and aliases them to `expected_schema`'s field names. +/// +/// [`ProjectionExec`] can rename fields. When the expected field is nullable and the input +/// field is not, this helper also widens nullability with a same-type [`CastExpr`]. It rejects +/// differences that projection cannot safely normalize exactly, such as data type, metadata, +/// schema metadata, and nullability narrowing. +pub fn project_plan_to_schema( + input: Arc<dyn ExecutionPlan>, + expected_schema: &SchemaRef, +) -> Result<Arc<dyn ExecutionPlan>> { + let input_schema = input.schema(); + if input_schema.as_ref() == expected_schema.as_ref() { + return Ok(input); + } + + if input_schema.fields().len() != expected_schema.fields().len() { + return plan_err!( + "Cannot project plan to expected schema: expected {} column(s), got {}", + expected_schema.fields().len(), + input_schema.fields().len() + ); + } + + if input_schema.metadata() != expected_schema.metadata() { + return plan_err!( + "Cannot project plan to expected schema: schema metadata differ" + ); + } + + if let Some((i, input_field, expected_field, mismatch)) = input_schema + .fields() + .iter() + .zip(expected_schema.fields().iter()) + .enumerate() + .find_map(|(i, (input_field, expected_field))| { + if input_field.data_type() != expected_field.data_type() { + Some((i, input_field, expected_field, "data type")) + } else if input_field.is_nullable() && !expected_field.is_nullable() { + Some((i, input_field, expected_field, "nullability")) + } else if input_field.metadata() != expected_field.metadata() { + Some((i, input_field, expected_field, "metadata")) + } else { + None + } + }) + { + return plan_err!( + "Cannot project plan column {i} ('{}') to expected output field '{}': \ + field {mismatch} differs (input field: {:?}, expected field: {:?})", + input_field.name(), + expected_field.name(), + input_field, + expected_field + ); + } + + let projection_exprs = expected_schema + .fields() + .iter() + .enumerate() + .map(|(i, expected_field)| { + let input_field = input_schema.field(i); + let column = Arc::new(Column::new(input_field.name(), i)); + let expr = if !input_field.is_nullable() && expected_field.is_nullable() { + Arc::new(CastExpr::new_with_target_field( + column, + Arc::clone(expected_field), + None, + )) as _ + } else { + column as _ + }; + ProjectionExpr { + expr, + alias: expected_field.name().clone(), + } + }) + .collect::<Vec<_>>(); + + let projection = ProjectionExec::try_new(projection_exprs, input)?; + debug_assert_eq!(projection.schema().as_ref(), expected_schema.as_ref()); + Ok(Arc::new(projection)) +} + +/// If running in a tokio context spawns the execution of `stream` to a separate task +/// allowing it to execute in parallel with an intermediate buffer of size `buffer`. +/// At most `buffer` record batches will be produced ahead of the consumer. +pub fn spawn_buffered( + mut input: SendableRecordBatchStream, + buffer: usize, +) -> SendableRecordBatchStream { + // Use tokio only if running from a multi-thread tokio context + match tokio::runtime::Handle::try_current() { + Ok(handle) + if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => + { + let mut builder = RecordBatchReceiverStream::builder(input.schema(), buffer); + + let sender = builder.tx(); + + builder.spawn(async move { + // We call `reserve` (which waits until there's room for at least 1 message in the + // channel buffer) **before** polling from input to ensure we hold a maximum of + // `buffer` record batches in memory. + // Polling from input and then calling send() would block when the channel is full + // so it would essentially hold `buffer` + 1 record batches: + // * `buffer`: this many elements would live inside the channel, since this is the + // channel's capacity + // * 1 extra RecordBatch which was produced, but there was no room for it in the + // channel, so it's being owned by the send() future, which keeps the batch in + // memory while it waits for a slot to free up + while let Ok(permit) = sender.reserve().await { + // Receiver dropped when query is shutdown early (e.g., limit) or error, + // no need to return propagate the send error. + match input.next().await { + Some(item) => permit.send(item), + None => break, + } + } + + Ok(()) + }); + + builder.build() + } + _ => input, + } +} + +/// Computes the statistics for an in-memory RecordBatch +/// +/// Only computes statistics that are in arrows metadata (num rows, byte size and nulls) +/// and does not apply any kernel on the actual data. +pub fn compute_record_batch_statistics( + batches: &[Vec<RecordBatch>], + schema: &Schema, + projection: Option<Vec<usize>>, +) -> Statistics { + let nb_rows = batches.iter().flatten().map(RecordBatch::num_rows).sum(); + + let projection = match projection { + Some(p) => p, + None => (0..schema.fields().len()).collect(), + }; + + let total_byte_size = batches + .iter() + .flatten() + .map(|b| { + projection + .iter() + .map(|index| b.column(*index).get_array_memory_size()) + .sum::<usize>() + }) + .sum(); + + let mut null_counts = vec![0; projection.len()]; + + for partition in batches.iter() { + for batch in partition { + for (stat_index, col_index) in projection.iter().enumerate() { + null_counts[stat_index] += batch + .column(*col_index) + .logical_nulls() + .map(|nulls| nulls.null_count()) + .unwrap_or_default(); + } + } + } + let column_statistics = null_counts + .into_iter() + .map(|null_count| { + let mut s = ColumnStatistics::new_unknown(); + s.null_count = Precision::Exact(null_count); + s + }) + .collect(); + + Statistics { + num_rows: Precision::Exact(nb_rows), + total_byte_size: Precision::Exact(total_byte_size), + column_statistics, + } +} + +/// Checks if the given projection is valid for the given schema. +pub fn can_project(schema: &SchemaRef, projection: Option<&[usize]>) -> Result<()> { + match projection { + Some(columns) => { + if columns + .iter() + .max() + .is_some_and(|&i| i >= schema.fields().len()) + { + Err(arrow::error::ArrowError::SchemaError(format!( + "project index {} out of bounds, max field {}", + columns.iter().max().unwrap(), + schema.fields().len() + )) + .into()) + } else { + Ok(()) + } + } + None => Ok(()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExec; + + use crate::stream::RecordBatchStreamAdapter; + use futures::stream; + use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use arrow::{ + array::{Float32Array, Float64Array, Int32Array, UInt64Array}, + datatypes::{DataType, Field, Schema}, + }; + + fn empty_exec(fields: Vec<Field>) -> Arc<dyn ExecutionPlan> { + Arc::new(EmptyExec::new(Arc::new(Schema::new(fields)))) + } + + #[test] + fn test_compute_record_batch_statistics_empty() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("f32", DataType::Float32, false), + Field::new("f64", DataType::Float64, false), + ])); + let stats = compute_record_batch_statistics(&[], &schema, Some(vec![0, 1])); + + assert_eq!(stats.num_rows, Precision::Exact(0)); + assert_eq!(stats.total_byte_size, Precision::Exact(0)); + Ok(()) + } + + #[test] + fn test_compute_record_batch_statistics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("f32", DataType::Float32, false), + Field::new("f64", DataType::Float64, false), + Field::new("u64", DataType::UInt64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Float32Array::from(vec![1., 2., 3.])), + Arc::new(Float64Array::from(vec![9., 8., 7.])), + Arc::new(UInt64Array::from(vec![4, 5, 6])), + ], + )?; + + // Just select f32,f64 + let select_projection = Some(vec![0, 1]); + let byte_size = batch + .project(&select_projection.clone().unwrap()) + .unwrap() + .get_array_memory_size(); + + let actual = + compute_record_batch_statistics(&[vec![batch]], &schema, select_projection); + + let expected = Statistics { + num_rows: Precision::Exact(3), + total_byte_size: Precision::Exact(byte_size), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Absent, + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Absent, + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }, + ], + }; + + assert_eq!(actual, expected); + Ok(()) + } + + #[test] + fn test_compute_record_batch_statistics_null() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("u64", DataType::UInt64, true)])); + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt64Array::from(vec![Some(1), None, None]))], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt64Array::from(vec![Some(1), Some(2), None]))], + )?; + let byte_size = batch1.get_array_memory_size() + batch2.get_array_memory_size(); + let actual = + compute_record_batch_statistics(&[vec![batch1], vec![batch2]], &schema, None); + + let expected = Statistics { + num_rows: Precision::Exact(6), + total_byte_size: Precision::Exact(byte_size), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Absent, + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + null_count: Precision::Exact(3), + byte_size: Precision::Absent, + }], + }; + + assert_eq!(actual, expected); + Ok(()) + } + + #[test] + fn project_plan_to_schema_returns_input_when_schema_matches() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let input: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let result = project_plan_to_schema(Arc::clone(&input), &schema)?; + + assert!(Arc::ptr_eq(&input, &result)); + Ok(()) + } + + #[test] + fn project_plan_to_schema_aliases_field_names_with_projection_exec() -> Result<()> { + let input = empty_exec(vec![ + Field::new("recursive_a", DataType::Int32, false), + Field::new("recursive_b", DataType::Utf8, true), + ]); + let expected_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ])); + + let result = project_plan_to_schema(Arc::clone(&input), &expected_schema)?; + + let projection = result + .downcast_ref::<ProjectionExec>() + .expect("schema rename should use ProjectionExec"); + assert!(Arc::ptr_eq(projection.input(), &input)); + assert_eq!(projection.schema(), expected_schema); + assert_eq!(projection.expr()[0].alias, "a"); + assert_eq!(projection.expr()[1].alias, "b"); + Ok(()) + } + + #[test] + fn project_plan_to_schema_preserves_matching_metadata_while_renaming() -> Result<()> { + let field_metadata = HashMap::from([("key".to_string(), "value".to_string())]); + let schema_metadata = + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]); + let input_schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("input", DataType::Int32, false) + .with_metadata(field_metadata.clone()), + ], + schema_metadata.clone(), + )); + let input: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(input_schema)); + let expected_schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("expected", DataType::Int32, false) + .with_metadata(field_metadata), + ], + schema_metadata, + )); + + let result = project_plan_to_schema(input, &expected_schema)?; + + assert_eq!(result.schema(), expected_schema); + Ok(()) + } + + #[test] + fn project_plan_to_schema_errors_on_column_count_mismatch() { + let input = empty_exec(vec![Field::new("a", DataType::Int32, false)]); + let expected_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); + assert!(err.to_string().contains("expected 2 column")); + } + + #[test] + fn project_plan_to_schema_errors_on_type_mismatch() { + let input = empty_exec(vec![Field::new("a", DataType::Int32, false)]); + let expected_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); + + let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); + assert!(err.to_string().contains("field data type differs")); + } + + #[test] + fn project_plan_to_schema_widens_nullability() -> Result<()> { + let input = empty_exec(vec![Field::new("a", DataType::Int32, false)]); + let expected_schema = Arc::new(Schema::new(vec![Field::new( + "renamed", + DataType::Int32, + true, + )])); + + let result = project_plan_to_schema(input, &expected_schema)?; + + assert_eq!(result.schema(), expected_schema); + Ok(()) + } + + #[test] + fn project_plan_to_schema_errors_on_nullability_narrowing() { + let input = empty_exec(vec![Field::new("a", DataType::Int32, true)]); + let expected_schema = Arc::new(Schema::new(vec![Field::new( + "renamed", + DataType::Int32, + false, + )])); + + let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); + assert!(err.to_string().contains("field nullability differs")); + } + + #[test] + fn project_plan_to_schema_errors_on_field_metadata_mismatch() { + let input = + empty_exec(vec![Field::new("a", DataType::Int32, false).with_metadata( + HashMap::from([("source".to_string(), "input".to_string())]), + )]); + let expected_schema = Arc::new(Schema::new(vec![ + Field::new("renamed", DataType::Int32, false).with_metadata(HashMap::from([ + ("source".to_string(), "expected".to_string()), + ])), + ])); + + let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); + assert!(err.to_string().contains("field metadata differs")); + } + + #[test] + fn project_plan_to_schema_errors_on_schema_metadata_mismatch() { + let input_schema = Arc::new(Schema::new_with_metadata( + vec![Field::new("a", DataType::Int32, false)], + HashMap::from([("source".to_string(), "input".to_string())]), + )); + let input: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(input_schema)); + let expected_schema = Arc::new(Schema::new_with_metadata( + vec![Field::new("renamed", DataType::Int32, false)], + HashMap::from([("source".to_string(), "expected".to_string())]), + )); + + let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); + assert!(err.to_string().contains("schema metadata differ")); + } + + /// Verifies that `spawn_buffered` holds exactly `buffer` record batches in memory + /// when no receiver is polling + async fn spawn_buffered_max_in_flight_batches(buffer_size: usize) { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let num_batches = 10; + + let produced_count = Arc::new(AtomicUsize::new(0)); + let produced_clone = Arc::clone(&produced_count); + let schema_clone = Arc::clone(&schema); + + // Stream increments the counter each time a batch is pulled by the producer. + let input_stream = stream::unfold(0usize, move |i| { + let schema = Arc::clone(&schema_clone); + let counter = Arc::clone(&produced_clone); + async move { + if i >= num_batches { + return None; + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![i as i32]))], + ) + .unwrap(); + counter.fetch_add(1, Ordering::SeqCst); + Some((Ok(batch), i + 1)) + } + }); + + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + input_stream, + )); + // Drop the returned stream immediately so no receiver is ever polled. + let _buffered = spawn_buffered(input, buffer_size); + + // Give the producer task time to fill the channel and stall on send(). + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + + assert_eq!( + produced_count.load(Ordering::SeqCst), + buffer_size, + "expected exactly {buffer_size} batch(es) in memory with no receiver polling" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_spawn_buffered_max_in_flight_batches() { + spawn_buffered_max_in_flight_batches(1).await; + spawn_buffered_max_in_flight_batches(2).await; + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/coop.rs b/native/vendor/datafusion-physical-plan/src/coop.rs new file mode 100644 index 0000000..9e27b26 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/coop.rs
@@ -0,0 +1,517 @@ +// 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. + +//! Utilities for improved cooperative scheduling. +//! +//! # Cooperative scheduling +//! +//! A single call to `poll_next` on a top-level [`Stream`] may potentially perform a lot of work +//! before it returns a `Poll::Pending`. Think for instance of calculating an aggregation over a +//! large dataset. +//! +//! If a `Stream` runs for a long period of time without yielding back to the Tokio executor, +//! it can starve other tasks waiting on that executor to execute them. +//! Additionally, this prevents the query execution from being cancelled. +//! +//! For more background, please also see the [Using Rust async for Query Execution and Cancelling Long-Running Queries blog] +//! +//! [Using Rust async for Query Execution and Cancelling Long-Running Queries blog]: https://datafusion.apache.org/blog/2025/06/30/cancellation +//! +//! To ensure that `Stream` implementations yield regularly, operators can insert explicit yield +//! points using the utilities in this module. For most operators this is **not** necessary. The +//! `Stream`s of the built-in DataFusion operators that generate (rather than manipulate) +//! `RecordBatch`es such as `DataSourceExec` and those that eagerly consume `RecordBatch`es +//! (for instance, `RepartitionExec`) contain yield points that will make most query `Stream`s yield +//! periodically. +//! +//! There are a couple of types of operators that _should_ insert yield points: +//! - New source operators that do not make use of Tokio resources +//! - Exchange like operators that do not use Tokio's `Channel` implementation to pass data between +//! tasks +//! +//! ## Adding yield points +//! +//! Yield points can be inserted manually using the facilities provided by the +//! [Tokio coop module](https://docs.rs/tokio/latest/tokio/task/coop/index.html) such as +//! [`tokio::task::coop::consume_budget`](https://docs.rs/tokio/latest/tokio/task/coop/fn.consume_budget.html). +//! +//! Another option is to use the wrapper `Stream` implementation provided by this module which will +//! consume a unit of task budget every time a `RecordBatch` is produced. +//! Wrapper `Stream`s can be created using the [`cooperative`] and [`make_cooperative`] functions. +//! +//! [`cooperative`] is a generic function that takes ownership of the wrapped [`RecordBatchStream`]. +//! This function has the benefit of not requiring an additional heap allocation and can avoid +//! dynamic dispatch. +//! +//! [`make_cooperative`] is a non-generic function that wraps a [`SendableRecordBatchStream`]. This +//! can be used to wrap dynamically typed, heap allocated [`RecordBatchStream`]s. +//! +//! ## Automatic cooperation +//! +//! The `EnsureCooperative` physical optimizer rule, which is included in the default set of +//! optimizer rules, inspects query plans for potential cooperative scheduling issues. +//! It injects the [`CooperativeExec`] wrapper `ExecutionPlan` into the query plan where necessary. +//! This `ExecutionPlan` uses [`make_cooperative`] to wrap the `Stream` of its input. +//! +//! The optimizer rule currently checks the plan for exchange-like operators and leave operators +//! that report [`SchedulingType::NonCooperative`] in their [plan properties](ExecutionPlan::properties). + +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_physical_expr::PhysicalExpr; +#[cfg(datafusion_coop = "tokio_fallback")] +use futures::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use crate::execution_plan::CardinalityEffect::{self, Equal}; +use crate::filter_pushdown::{ + ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + SortOrderPushdownResult, validate_child_count, +}; +use arrow::record_batch::RecordBatch; +use arrow_schema::Schema; +use datafusion_common::{Result, Statistics}; +use datafusion_execution::TaskContext; + +use crate::execution_plan::{SchedulingType, replace_children_if_necessary}; +use crate::stream::RecordBatchStreamAdapter; +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use futures::{Stream, StreamExt}; + +/// A stream that passes record batches through unchanged while cooperating with the Tokio runtime. +/// It consumes cooperative scheduling budget for each returned [`RecordBatch`], +/// allowing other tasks to execute when the budget is exhausted. +/// +/// See the [module level documentation](crate::coop) for an in-depth discussion. +pub struct CooperativeStream<T> +where + T: RecordBatchStream + Unpin, +{ + inner: T, + #[cfg(datafusion_coop = "per_stream")] + budget: u8, +} + +#[cfg(datafusion_coop = "per_stream")] +// Magic value that matches Tokio's task budget value +const YIELD_FREQUENCY: u8 = 128; + +impl<T> CooperativeStream<T> +where + T: RecordBatchStream + Unpin, +{ + /// Creates a new `CooperativeStream` that wraps the provided stream. + /// The resulting stream will cooperate with the Tokio scheduler by consuming a unit of + /// scheduling budget when the wrapped `Stream` returns a record batch. + pub fn new(inner: T) -> Self { + Self { + inner, + #[cfg(datafusion_coop = "per_stream")] + budget: YIELD_FREQUENCY, + } + } +} + +impl<T> Stream for CooperativeStream<T> +where + T: RecordBatchStream + Unpin, +{ + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + #[cfg(any( + datafusion_coop = "tokio", + not(any( + datafusion_coop = "tokio_fallback", + datafusion_coop = "per_stream" + )) + ))] + { + let coop = std::task::ready!(tokio::task::coop::poll_proceed(cx)); + let value = self.inner.poll_next_unpin(cx); + if value.is_ready() { + coop.made_progress(); + } + value + } + + #[cfg(datafusion_coop = "tokio_fallback")] + { + // This is a temporary placeholder implementation that may have slightly + // worse performance compared to `poll_proceed` + if !tokio::task::coop::has_budget_remaining() { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + + let value = self.inner.poll_next_unpin(cx); + if value.is_ready() { + // In contrast to `poll_proceed` we are not able to consume + // budget before proceeding to do work. Instead, we try to consume budget + // after the work has been done and just assume that that succeeded. + // The poll result is ignored because we don't want to discard + // or buffer the Ready result we got from the inner stream. + let consume = tokio::task::coop::consume_budget(); + let consume_ref = std::pin::pin!(consume); + let _ = consume_ref.poll(cx); + } + value + } + + #[cfg(datafusion_coop = "per_stream")] + { + if self.budget == 0 { + self.budget = YIELD_FREQUENCY; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + + let value = { self.inner.poll_next_unpin(cx) }; + + if value.is_ready() { + self.budget -= 1; + } else { + self.budget = YIELD_FREQUENCY; + } + value + } + } +} + +impl<T> RecordBatchStream for CooperativeStream<T> +where + T: RecordBatchStream + Unpin, +{ + fn schema(&self) -> Arc<Schema> { + self.inner.schema() + } +} + +/// An execution plan decorator that enables cooperative multitasking. +/// It wraps the streams produced by its input execution plan using the [`make_cooperative`] function, +/// which makes the stream participate in Tokio cooperative scheduling. +#[derive(Debug, Clone)] +pub struct CooperativeExec { + input: Arc<dyn ExecutionPlan>, + properties: Arc<PlanProperties>, +} + +impl CooperativeExec { + /// Creates a new `CooperativeExec` operator that wraps the given input execution plan. + pub fn new(input: Arc<dyn ExecutionPlan>) -> Self { + let properties = PlanProperties::clone(input.properties()) + .with_scheduling_type(SchedulingType::Cooperative) + .into(); + + Self { input, properties } + } + + /// Returns a reference to the wrapped input execution plan. + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } +} + +impl DisplayAs for CooperativeExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter<'_>, + ) -> std::fmt::Result { + write!(f, "CooperativeExec") + } +} + +impl ExecutionPlan for CooperativeExec { + fn name(&self) -> &str { + "CooperativeExec" + } + + fn schema(&self) -> Arc<Schema> { + self.input.schema() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.properties + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true; self.children().len()] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + task_ctx: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let child_stream = self.input.execute(partition, task_ctx)?; + Ok(make_cooperative(child_stream)) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::clone(&input_stats[0])) + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + fn cardinality_effect(&self) -> CardinalityEffect { + Equal + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + match self.input.try_swapping_with_projection(projection)? { + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), + None => Ok(None), + } + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + FilterDescription::from_children(parent_filters, &self.children()) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + Ok(FilterPushdownPropagation::if_all(child_pushdown_result)) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + let child = self.input(); + + match child.try_pushdown_sort(order)? { + SortOrderPushdownResult::Exact { inner } => { + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; + Ok(SortOrderPushdownResult::Exact { inner: new_exec }) + } + SortOrderPushdownResult::Inexact { inner } => { + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; + Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) + } + SortOrderPushdownResult::Unsupported => { + Ok(SortOrderPushdownResult::Unsupported) + } + } + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Cooperative(Box::new( + protobuf::CooperativeExecNode { + input: Some(Box::new(input)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CooperativeExec { + /// Reconstruct a [`CooperativeExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let cooperative = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Cooperative, + "CooperativeExec", + ); + let input = ctx.decode_required_child( + cooperative.input.as_deref(), + "CooperativeExec", + "input", + )?; + Ok(Arc::new(CooperativeExec::new(input))) + } +} + +/// Creates a [`CooperativeStream`] wrapper around the given [`RecordBatchStream`]. +/// This wrapper collaborates with the Tokio cooperative scheduler by consuming a unit of +/// scheduling budget for each returned record batch. +pub fn cooperative<T>(stream: T) -> CooperativeStream<T> +where + T: RecordBatchStream + Unpin + Send + 'static, +{ + CooperativeStream::new(stream) +} + +/// Wraps a `SendableRecordBatchStream` inside a [`CooperativeStream`] to enable cooperative multitasking. +/// Since `SendableRecordBatchStream` is a `dyn RecordBatchStream` this requires the use of dynamic +/// method dispatch. +/// When the stream type is statically known, consider use the generic [`cooperative`] function +/// to allow static method dispatch. +pub fn make_cooperative(stream: SendableRecordBatchStream) -> SendableRecordBatchStream { + // TODO is there a more elegant way to overload cooperative + Box::pin(cooperative(RecordBatchStreamAdapter::new( + stream.schema(), + stream, + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_schema::SchemaRef; + + use futures::stream; + + // This is the hardcoded value Tokio uses + const TASK_BUDGET: usize = 128; + + /// Helper: construct a SendableRecordBatchStream containing `n` empty batches + fn make_empty_batches(n: usize) -> SendableRecordBatchStream { + let schema: SchemaRef = Arc::new(Schema::empty()); + let schema_for_stream = Arc::clone(&schema); + + let s = + stream::iter((0..n).map(move |_| { + Ok(RecordBatch::new_empty(Arc::clone(&schema_for_stream))) + })); + + Box::pin(RecordBatchStreamAdapter::new(schema, s)) + } + + #[tokio::test] + async fn yield_less_than_threshold() -> Result<()> { + let count = TASK_BUDGET - 10; + let inner = make_empty_batches(count); + let out = make_cooperative(inner).collect::<Vec<_>>().await; + assert_eq!(out.len(), count); + Ok(()) + } + + #[tokio::test] + async fn yield_equal_to_threshold() -> Result<()> { + let count = TASK_BUDGET; + let inner = make_empty_batches(count); + let out = make_cooperative(inner).collect::<Vec<_>>().await; + assert_eq!(out.len(), count); + Ok(()) + } + + #[tokio::test] + async fn yield_more_than_threshold() -> Result<()> { + let count = TASK_BUDGET + 20; + let inner = make_empty_batches(count); + let out = make_cooperative(inner).collect::<Vec<_>>().await; + assert_eq!(out.len(), count); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/display.rs b/native/vendor/datafusion-physical-plan/src/display.rs new file mode 100644 index 0000000..d2bdcef --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/display.rs
@@ -0,0 +1,1886 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Implementation of physical plan display. See +//! [`crate::displayable`] for examples of how to format + +use std::collections::{BTreeMap, HashMap}; +use std::fmt; +use std::fmt::Formatter; +use std::time::Duration; + +use arrow::datatypes::SchemaRef; + +use datafusion_common::display::{GraphvizBuilder, PlanType, StringifiedPlan}; +use datafusion_expr::display_schema; +use datafusion_physical_expr::LexOrdering; + +use crate::metrics::{MetricCategory, MetricType, MetricValue}; +use crate::render_tree::RenderTree; + +use crate::statistics::{StatisticsArgs, StatisticsContext}; + +use super::{ExecutionPlan, ExecutionPlanVisitor, accept}; + +/// Options for controlling how each [`ExecutionPlan`] should format itself +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum DisplayFormatType { + /// Default, compact format. Example: `FilterExec: c12 < 10.0` + /// + /// This format is designed to provide a detailed textual description + /// of all parts of the plan. + Default, + /// Verbose, showing all available details. + /// + /// This form is even more detailed than [`Self::Default`] + Verbose, + /// TreeRender, displayed in the `tree` explain type. + /// + /// This format is inspired by DuckDB's explain plans. The information + /// presented should be "user friendly", and contain only the most relevant + /// information for understanding a plan. It should NOT contain the same level + /// of detail information as the [`Self::Default`] format. + /// + /// In this mode, each line has one of two formats: + /// + /// 1. A string without a `=`, which is printed in its own line + /// + /// 2. A string with a `=` that is treated as a `key=value pair`. Everything + /// before the first `=` is treated as the key, and everything after the + /// first `=` is treated as the value. + /// + /// For example, if the output of `TreeRender` is this: + /// ```text + /// Parquet + /// partition_sizes=[1] + /// ``` + /// + /// It is rendered in the center of a box in the following way: + /// + /// ```text + /// ┌───────────────────────────┐ + /// │ DataSourceExec │ + /// │ -------------------- │ + /// │ partition_sizes: [1] │ + /// │ Parquet │ + /// └───────────────────────────┘ + /// ``` + TreeRender, +} + +/// Wraps an `ExecutionPlan` with various methods for formatting +/// +/// +/// # Example +/// ``` +/// # use std::sync::Arc; +/// # use arrow::datatypes::{Field, Schema, DataType}; +/// # use datafusion_expr::Operator; +/// # use datafusion_physical_expr::expressions::{binary, col, lit}; +/// # use datafusion_physical_plan::{displayable, ExecutionPlan}; +/// # use datafusion_physical_plan::empty::EmptyExec; +/// # use datafusion_physical_plan::filter::FilterExec; +/// # let schema = Schema::new(vec![Field::new("i", DataType::Int32, false)]); +/// # let plan = EmptyExec::new(Arc::new(schema)); +/// # let i = col("i", &plan.schema()).unwrap(); +/// # let predicate = binary(i, Operator::Eq, lit(1), &plan.schema()).unwrap(); +/// # let plan: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new(predicate, Arc::new(plan)).unwrap()); +/// // Get a one line description (Displayable) +/// let display_plan = displayable(plan.as_ref()); +/// +/// // you can use the returned objects to format plans +/// // where you can use `Display` such as format! or println! +/// assert_eq!( +/// &format!("The plan is: {}", display_plan.one_line()), +/// "The plan is: FilterExec: i@0 = 1\n" +/// ); +/// // You can also print out the plan and its children in indented mode +/// assert_eq!(display_plan.indent(false).to_string(), +/// "FilterExec: i@0 = 1\ +/// \n EmptyExec\ +/// \n" +/// ); +/// ``` +#[derive(Debug, Clone)] +pub struct DisplayableExecutionPlan<'a> { + inner: &'a dyn ExecutionPlan, + /// How to show metrics + show_metrics: ShowMetrics, + /// If statistics should be displayed + show_statistics: bool, + /// If schema should be displayed. See [`Self::set_show_schema`] + show_schema: bool, + /// Which metric categories should be included when rendering + metric_types: Vec<MetricType>, + /// Optional filter by semantic category (rows / bytes / timing). + /// `None` means show all categories; `Some(vec![])` means plan-only. + metric_categories: Option<Vec<MetricCategory>>, + /// Optional filter by metric names. Only metric names in this list + /// will be rendered. + metric_names: Option<Vec<String>>, + // (TreeRender) Maximum total width of the rendered tree + tree_maximum_render_width: usize, + /// Optional summary totals (currently only used by `pgjson`) — the total + /// row count and wall-clock duration of the `AnalyzeExec` execution. + summary: Option<AnalyzeSummary>, +} + +/// Summary information attached to the root of an `EXPLAIN ANALYZE` +/// pgjson render. +#[derive(Debug, Clone, Copy)] +struct AnalyzeSummary { + total_rows: Option<usize>, + duration: Option<Duration>, +} + +impl<'a> DisplayableExecutionPlan<'a> { + fn default_metric_types() -> Vec<MetricType> { + vec![MetricType::Summary, MetricType::Dev] + } + + /// Create a wrapper around an [`ExecutionPlan`] which can be + /// pretty printed in a variety of ways + pub fn new(inner: &'a dyn ExecutionPlan) -> Self { + Self { + inner, + show_metrics: ShowMetrics::None, + show_statistics: false, + show_schema: false, + metric_types: Self::default_metric_types(), + metric_categories: None, + metric_names: None, + tree_maximum_render_width: 240, + summary: None, + } + } + + /// Create a wrapper around an [`ExecutionPlan`] which can be + /// pretty printed in a variety of ways that also shows aggregated + /// metrics + pub fn with_metrics(inner: &'a dyn ExecutionPlan) -> Self { + Self { + inner, + show_metrics: ShowMetrics::Aggregated, + show_statistics: false, + show_schema: false, + metric_types: Self::default_metric_types(), + metric_categories: None, + metric_names: None, + tree_maximum_render_width: 240, + summary: None, + } + } + + /// Create a wrapper around an [`ExecutionPlan`] which can be + /// pretty printed in a variety of ways that also shows all low + /// level metrics + pub fn with_full_metrics(inner: &'a dyn ExecutionPlan) -> Self { + Self { + inner, + show_metrics: ShowMetrics::Full, + show_statistics: false, + show_schema: false, + metric_types: Self::default_metric_types(), + metric_categories: None, + metric_names: None, + tree_maximum_render_width: 240, + summary: None, + } + } + + /// Enable display of schema + /// + /// If true, plans will be displayed with schema information at the end + /// of each line. The format is `schema=[[a:Int32;N, b:Int32;N, c:Int32;N]]` + pub fn set_show_schema(mut self, show_schema: bool) -> Self { + self.show_schema = show_schema; + self + } + + /// Enable display of statistics + pub fn set_show_statistics(mut self, show_statistics: bool) -> Self { + self.show_statistics = show_statistics; + self + } + + /// Specify which metric types should be rendered alongside the plan + pub fn set_metric_types(mut self, metric_types: Vec<MetricType>) -> Self { + self.metric_types = metric_types; + self + } + + /// Specify which metric categories to include. + /// + /// - `None` means show all categories (default). + /// - `Some(vec![])` means plan-only — suppress all metrics. + /// - `Some(vec![Rows])` means show only row-count metrics (plus + /// uncategorized metrics). + /// + /// See [`MetricCategory`] for the determinism properties of each + /// category. + pub fn set_metric_categories( + mut self, + metric_categories: Option<Vec<MetricCategory>>, + ) -> Self { + self.metric_categories = metric_categories; + self + } + + /// Specify which metric names to include. + /// + /// - An empty vector means plan-only — suppress all metrics. + /// - `vec!["metric_1"]` means show only the metric named `metric_1`. + /// + /// Name filtering is intersected with other types of filters, like metric + /// category and metric type. + pub fn set_metric_names(mut self, metric_names: Vec<String>) -> Self { + self.metric_names = Some(metric_names); + self + } + + /// Set the maximum render width for the tree format + pub fn set_tree_maximum_render_width(mut self, width: usize) -> Self { + self.tree_maximum_render_width = width; + self + } + + /// Attach an `EXPLAIN ANALYZE` summary (total output rows and duration) + /// to the rendered output. Currently only used by [`Self::pgjson`], which + /// serializes the summary alongside the root plan object. + pub fn set_summary( + mut self, + total_rows: Option<usize>, + duration: Option<Duration>, + ) -> Self { + self.summary = Some(AnalyzeSummary { + total_rows, + duration, + }); + self + } + + /// Return a `format`able structure that produces a single line + /// per node. + /// + /// ```text + /// ProjectionExec: expr=[a] + /// CoalesceBatchesExec: target_batch_size=8192 + /// FilterExec: a < 5 + /// RepartitionExec: partitioning=RoundRobinBatch(16) + /// DataSourceExec: source=...", + /// ``` + pub fn indent(&self, verbose: bool) -> impl fmt::Display + 'a { + let format_type = if verbose { + DisplayFormatType::Verbose + } else { + DisplayFormatType::Default + }; + struct Wrapper<'a> { + format_type: DisplayFormatType, + plan: &'a dyn ExecutionPlan, + show_metrics: ShowMetrics, + show_statistics: bool, + show_schema: bool, + metric_types: Vec<MetricType>, + metric_categories: Option<Vec<MetricCategory>>, + metric_names: Option<Vec<String>>, + } + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let mut visitor = IndentVisitor { + t: self.format_type, + f, + indent: 0, + show_metrics: self.show_metrics, + show_statistics: self.show_statistics, + show_schema: self.show_schema, + metric_types: &self.metric_types, + metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), + }; + accept(self.plan, &mut visitor) + } + } + Wrapper { + format_type, + plan: self.inner, + show_metrics: self.show_metrics, + show_statistics: self.show_statistics, + show_schema: self.show_schema, + metric_types: self.metric_types.clone(), + metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), + } + } + + /// Returns a `format`able structure that produces graphviz format for execution plan, which can + /// be directly visualized [here](https://dreampuf.github.io/GraphvizOnline). + /// + /// An example is + /// ```dot + /// strict digraph dot_plan { + // 0[label="ProjectionExec: expr=[id@0 + 2 as employee.id + Int32(2)]",tooltip=""] + // 1[label="EmptyExec",tooltip=""] + // 0 -> 1 + // } + /// ``` + pub fn graphviz(&self) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + show_metrics: ShowMetrics, + show_statistics: bool, + metric_types: Vec<MetricType>, + metric_categories: Option<Vec<MetricCategory>>, + metric_names: Option<Vec<String>>, + } + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let t = DisplayFormatType::Default; + + let mut visitor = GraphvizVisitor { + f, + t, + show_metrics: self.show_metrics, + show_statistics: self.show_statistics, + metric_types: &self.metric_types, + metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), + graphviz_builder: GraphvizBuilder::default(), + parents: Vec::new(), + }; + + visitor.start_graph()?; + + accept(self.plan, &mut visitor)?; + + visitor.end_graph()?; + Ok(()) + } + } + + Wrapper { + plan: self.inner, + show_metrics: self.show_metrics, + show_statistics: self.show_statistics, + metric_types: self.metric_types.clone(), + metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), + } + } + + /// Formats the plan using a ASCII art like tree + /// + /// See [`DisplayFormatType::TreeRender`] for more details. + pub fn tree_render(&self) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + maximum_render_width: usize, + } + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let mut visitor = TreeRenderVisitor { + f, + maximum_render_width: self.maximum_render_width, + }; + visitor.visit(self.plan) + } + } + Wrapper { + plan: self.inner, + maximum_render_width: self.tree_maximum_render_width, + } + } + + /// Returns a `format`able structure that produces PostgreSQL-style JSON + /// output, mirroring the logical-plan pgjson format. + /// + /// Each node is rendered as a JSON object with: + /// - `"Node Type"` — `ExecutionPlan::name()` + /// - `"Details"` — the one-line `DisplayAs::Default` rendering + /// - `"Output"` — schema column names (when `set_show_schema(true)`) + /// - `"Actual Rows"` / `"Actual Total Time"` — PG-canonical metric keys + /// populated from `output_rows` / `elapsed_compute` when available + /// - `"Extras"` — remaining metrics keyed by DataFusion metric name + /// - `"Plans"` — array of child nodes + /// + /// When a summary has been set via [`Self::set_summary`], `"Total Rows"` + /// and `"Duration"` fields are attached at the root. + pub fn pgjson(&self, verbose: bool) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + verbose: bool, + show_metrics: ShowMetrics, + show_schema: bool, + metric_types: Vec<MetricType>, + metric_categories: Option<Vec<MetricCategory>>, + metric_names: Option<Vec<String>>, + summary: Option<AnalyzeSummary>, + } + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let mut visitor = PgJsonExecutionPlanVisitor { + verbose: self.verbose, + show_metrics: self.show_metrics, + show_schema: self.show_schema, + metric_types: &self.metric_types, + metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), + objects: HashMap::new(), + parent_ids: Vec::new(), + next_id: 0, + root: None, + }; + accept(self.plan, &mut visitor).map_err(|_| fmt::Error)?; + let root = visitor.root.ok_or(fmt::Error)?; + let mut root_entry = serde_json::json!({ "Plan": root }); + if let Some(summary) = self.summary { + if let Some(total_rows) = summary.total_rows { + root_entry["Total Rows"] = serde_json::Value::from(total_rows); + } + if let Some(duration) = summary.duration { + root_entry["Duration"] = + serde_json::Value::from(format!("{duration:?}")); + } + } + let doc = serde_json::Value::Array(vec![root_entry]); + write!( + f, + "{}", + serde_json::to_string_pretty(&doc).map_err(|_| fmt::Error)? + ) + } + } + + Wrapper { + plan: self.inner, + verbose, + show_metrics: self.show_metrics, + show_schema: self.show_schema, + metric_types: self.metric_types.clone(), + metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), + summary: self.summary, + } + } + + /// Return a single-line summary of the root of the plan + /// Example: `ProjectionExec: expr=[a@0 as a]`. + pub fn one_line(&self) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + show_metrics: ShowMetrics, + show_statistics: bool, + show_schema: bool, + metric_types: Vec<MetricType>, + metric_categories: Option<Vec<MetricCategory>>, + metric_names: Option<Vec<String>>, + } + + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let mut visitor = IndentVisitor { + f, + t: DisplayFormatType::Default, + indent: 0, + show_metrics: self.show_metrics, + show_statistics: self.show_statistics, + show_schema: self.show_schema, + metric_types: &self.metric_types, + metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), + }; + visitor.pre_visit(self.plan)?; + Ok(()) + } + } + + Wrapper { + plan: self.inner, + show_metrics: self.show_metrics, + show_statistics: self.show_statistics, + show_schema: self.show_schema, + metric_types: self.metric_types.clone(), + metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), + } + } + + #[deprecated(since = "47.0.0", note = "indent() or tree_render() instead")] + pub fn to_stringified( + &self, + verbose: bool, + plan_type: PlanType, + explain_format: DisplayFormatType, + ) -> StringifiedPlan { + match (&explain_format, &plan_type) { + (DisplayFormatType::TreeRender, PlanType::FinalPhysicalPlan) => { + StringifiedPlan::new(plan_type, self.tree_render().to_string()) + } + _ => StringifiedPlan::new(plan_type, self.indent(verbose).to_string()), + } + } +} + +/// Enum representing the different levels of metrics to display +#[derive(Debug, Clone, Copy)] +enum ShowMetrics { + /// Do not show any metrics + None, + + /// Show aggregated metrics across partition + Aggregated, + + /// Show full per-partition metrics + Full, +} + +/// Formats plans with a single line per node. +/// +/// # Example +/// +/// ```text +/// ProjectionExec: expr=[column1@0 + 2 as column1 + Int64(2)] +/// FilterExec: column1@0 = 5 +/// ValuesExec +/// ``` +struct IndentVisitor<'a, 'b> { + /// How to format each node + t: DisplayFormatType, + /// Write to this formatter + f: &'a mut Formatter<'b>, + /// Indent size + indent: usize, + /// How to show metrics + show_metrics: ShowMetrics, + /// If statistics should be displayed + show_statistics: bool, + /// If schema should be displayed + show_schema: bool, + /// Which metric types should be rendered + metric_types: &'a [MetricType], + /// Optional filter by semantic category (rows / bytes / timing). + metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, +} + +impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { + type Error = fmt::Error; + fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + write!(self.f, "{:indent$}", "", indent = self.indent * 2)?; + plan.fmt_as(self.t, self.f)?; + match self.show_metrics { + ShowMetrics::None => {} + ShowMetrics::Aggregated => { + if let Some(metrics) = plan.metrics() { + let mut metrics = metrics + .filter_by_metric_types(self.metric_types) + .aggregate_by_name() + .sorted_for_display() + .timestamps_removed(); + if let Some(cats) = self.metric_categories { + metrics = metrics.filter_by_categories(cats); + } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } + write!(self.f, ", metrics=[{metrics}]")?; + } else { + write!(self.f, ", metrics=[]")?; + } + } + ShowMetrics::Full => { + if let Some(metrics) = plan.metrics() { + let mut metrics = metrics.filter_by_metric_types(self.metric_types); + if let Some(cats) = self.metric_categories { + metrics = metrics.filter_by_categories(cats); + } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } + write!(self.f, ", metrics=[{metrics}]")?; + } else { + write!(self.f, ", metrics=[]")?; + } + } + } + if self.show_statistics { + let stats = StatisticsContext::new() + .compute(plan, &StatisticsArgs::new()) + .map_err(|_e| fmt::Error)?; + write!(self.f, ", statistics=[{stats}]")?; + } + if self.show_schema { + write!( + self.f, + ", schema={}", + display_schema(plan.schema().as_ref()) + )?; + } + writeln!(self.f)?; + self.indent += 1; + Ok(true) + } + + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + self.indent -= 1; + Ok(true) + } +} + +struct GraphvizVisitor<'a, 'b> { + f: &'a mut Formatter<'b>, + /// How to format each node + t: DisplayFormatType, + /// How to show metrics + show_metrics: ShowMetrics, + /// If statistics should be displayed + show_statistics: bool, + /// Which metric types should be rendered + metric_types: &'a [MetricType], + /// Optional filter by semantic category + metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, + + graphviz_builder: GraphvizBuilder, + /// Used to record parent node ids when visiting a plan. + parents: Vec<usize>, +} + +impl GraphvizVisitor<'_, '_> { + fn start_graph(&mut self) -> fmt::Result { + self.graphviz_builder.start_graph(self.f) + } + + fn end_graph(&mut self) -> fmt::Result { + self.graphviz_builder.end_graph(self.f) + } +} + +impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { + type Error = fmt::Error; + + fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + let id = self.graphviz_builder.next_id(); + + struct Wrapper<'a>(&'a dyn ExecutionPlan, DisplayFormatType); + + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt_as(self.1, f) + } + } + + let label = { format!("{}", Wrapper(plan, self.t)) }; + + let metrics = match self.show_metrics { + ShowMetrics::None => "".to_string(), + ShowMetrics::Aggregated => { + if let Some(metrics) = plan.metrics() { + let mut metrics = metrics + .filter_by_metric_types(self.metric_types) + .aggregate_by_name() + .sorted_for_display() + .timestamps_removed(); + if let Some(cats) = self.metric_categories { + metrics = metrics.filter_by_categories(cats); + } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } + format!("metrics=[{metrics}]") + } else { + "metrics=[]".to_string() + } + } + ShowMetrics::Full => { + if let Some(metrics) = plan.metrics() { + let mut metrics = metrics.filter_by_metric_types(self.metric_types); + if let Some(cats) = self.metric_categories { + metrics = metrics.filter_by_categories(cats); + } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } + format!("metrics=[{metrics}]") + } else { + "metrics=[]".to_string() + } + } + }; + + let statistics = if self.show_statistics { + let stats = StatisticsContext::new() + .compute(plan, &StatisticsArgs::new()) + .map_err(|_e| fmt::Error)?; + format!("statistics=[{stats}]") + } else { + "".to_string() + }; + + let delimiter = if !metrics.is_empty() && !statistics.is_empty() { + ", " + } else { + "" + }; + + self.graphviz_builder.add_node( + self.f, + id, + &label, + Some(&format!("{metrics}{delimiter}{statistics}")), + )?; + + if let Some(parent_node_id) = self.parents.last() { + self.graphviz_builder + .add_edge(self.f, *parent_node_id, id)?; + } + + self.parents.push(id); + + Ok(true) + } + + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + self.parents.pop(); + Ok(true) + } +} + +/// Formats physical plans into PostgreSQL-style JSON output with live +/// per-operator metrics. +/// +/// This visitor mirrors the logical-plan `PgJsonVisitor` in +/// `datafusion-expr`: during `pre_visit` it assembles a JSON object for the +/// current node; during `post_visit` it attaches that object into its +/// parent's `"Plans"` array (or stores it as the root). +struct PgJsonExecutionPlanVisitor<'a> { + verbose: bool, + show_metrics: ShowMetrics, + show_schema: bool, + metric_types: &'a [MetricType], + metric_categories: Option<&'a [MetricCategory]>, + metric_names: Option<&'a [String]>, + objects: HashMap<u32, serde_json::Value>, + parent_ids: Vec<u32>, + next_id: u32, + root: Option<serde_json::Value>, +} + +impl PgJsonExecutionPlanVisitor<'_> { + /// Produce the one-line `DisplayAs::Default` rendering of a node. + fn one_line_details(plan: &dyn ExecutionPlan) -> String { + struct One<'b>(&'b dyn ExecutionPlan); + impl fmt::Display for One<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } + } + // Some operators include internal newlines; collapse them so the + // rendered JSON value stays on a single line. + format!("{}", One(plan)) + .replace('\n', " ") + .trim() + .to_string() + } + + /// Render the given `MetricValue` into the most natural `serde_json::Value` + /// we can produce: a number for simple counts/gauges/times, a float-ms for + /// `ElapsedCompute`, and a string fallback for anything else. + fn metric_value_to_json(value: &MetricValue) -> serde_json::Value { + match value { + MetricValue::OutputRows(c) => serde_json::Value::from(c.value()), + MetricValue::SpillCount(c) + | MetricValue::OutputBatches(c) + | MetricValue::SpilledRows(c) => serde_json::Value::from(c.value()), + MetricValue::SpilledBytes(c) | MetricValue::OutputBytes(c) => { + serde_json::Value::from(c.value()) + } + MetricValue::CurrentMemoryUsage(g) => serde_json::Value::from(g.value()), + MetricValue::ElapsedCompute(t) => { + // Emit as float milliseconds to align with PG's + // `"Actual Total Time"` convention. DataFusion tracks compute + // time (summed across partitions), not wall time — visualizers + // should be read with that caveat in mind. + let ms = (t.value() as f64) / 1_000_000.0; + serde_json::Value::from(ms) + } + MetricValue::Count { count, .. } => serde_json::Value::from(count.value()), + MetricValue::Gauge { gauge, .. } => serde_json::Value::from(gauge.value()), + MetricValue::PeakMemoryUsage { gauge, .. } => { + serde_json::Value::from(gauge.value()) + } + MetricValue::Time { time, .. } => { + let ms = (time.value() as f64) / 1_000_000.0; + serde_json::Value::from(ms) + } + // Timestamps, PruningMetrics, Ratio, Custom: fall back to Display. + other => serde_json::Value::String(format!("{other}")), + } + } + + /// Populate `"Actual Rows"`, `"Actual Total Time"`, and `"Extras"` for + /// the given node from its aggregated `MetricsSet`, honoring the same + /// filtering pipeline used by `IndentVisitor`. + fn attach_metrics(&self, plan: &dyn ExecutionPlan, object: &mut serde_json::Value) { + if matches!(self.show_metrics, ShowMetrics::None) { + return; + } + let Some(metrics) = plan.metrics() else { + return; + }; + + let metrics = match self.show_metrics { + ShowMetrics::None => return, + ShowMetrics::Aggregated => metrics + .filter_by_metric_types(self.metric_types) + .aggregate_by_name() + .sorted_for_display() + .timestamps_removed(), + ShowMetrics::Full => metrics.filter_by_metric_types(self.metric_types), + }; + let metrics = if let Some(cats) = self.metric_categories { + metrics.filter_by_categories(cats) + } else { + metrics + }; + + let metrics = if let Some(names) = self.metric_names { + metrics.filter_by_names(names) + } else { + metrics + }; + + // Build the Extras bucket, while extracting PG-canonical keys to the + // top level. + let mut extras = serde_json::Map::new(); + for metric in metrics.iter() { + let value = metric.value(); + match value { + MetricValue::OutputRows(c) => { + object["Actual Rows"] = serde_json::Value::from(c.value()); + } + MetricValue::ElapsedCompute(t) => { + let ms = (t.value() as f64) / 1_000_000.0; + object["Actual Total Time"] = serde_json::Value::from(ms); + } + _ => { + extras.insert( + value.name().to_string(), + Self::metric_value_to_json(value), + ); + } + } + } + if !extras.is_empty() { + object["Extras"] = serde_json::Value::Object(extras); + } + } +} + +impl ExecutionPlanVisitor for PgJsonExecutionPlanVisitor<'_> { + type Error = fmt::Error; + + fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + let id = self.next_id; + self.next_id += 1; + + // Build fields in reading order: Node Type, Details, (schema), + // (metrics), Plans last — so the JSON output reads top-down like a + // PostgreSQL plan. + let mut object = serde_json::json!({ + "Node Type": plan.name(), + "Details": Self::one_line_details(plan), + }); + + if self.show_schema || self.verbose { + // Always include output columns when a caller asked for schema; + // also include them in verbose mode so the pgjson output mirrors + // the extra context shown by indent's verbose flag. + let columns: Vec<serde_json::Value> = plan + .schema() + .fields() + .iter() + .map(|f| serde_json::Value::String(f.name().to_string())) + .collect(); + object["Output"] = serde_json::Value::Array(columns); + } + + self.attach_metrics(plan, &mut object); + + object["Plans"] = serde_json::Value::Array(vec![]); + + self.objects.insert(id, object); + self.parent_ids.push(id); + Ok(true) + } + + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + let id = self.parent_ids.pop().ok_or(fmt::Error)?; + let current = self.objects.remove(&id).ok_or(fmt::Error)?; + + if let Some(parent_id) = self.parent_ids.last() { + let parent = self.objects.get_mut(parent_id).ok_or(fmt::Error)?; + let plans = parent + .get_mut("Plans") + .and_then(|p| p.as_array_mut()) + .ok_or(fmt::Error)?; + plans.push(current); + } else { + self.root = Some(current); + } + Ok(true) + } +} + +/// This module implements a tree-like art renderer for execution plans, +/// based on DuckDB's implementation: +/// <https://github.com/duckdb/duckdb/blob/main/src/include/duckdb/common/tree_renderer/text_tree_renderer.hpp> +/// +/// The rendered output looks like this: +/// ```text +/// ┌───────────────────────────┐ +/// │ CoalesceBatchesExec │ +/// └─────────────┬─────────────┘ +/// ┌─────────────┴─────────────┐ +/// │ HashJoinExec ├──────────────┐ +/// └─────────────┬─────────────┘ │ +/// ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +/// │ DataSourceExec ││ DataSourceExec │ +/// └───────────────────────────┘└───────────────────────────┘ +/// ``` +/// +/// The renderer uses a three-layer approach for each node: +/// 1. Top layer: renders the top borders and connections +/// 2. Content layer: renders the node content and vertical connections +/// 3. Bottom layer: renders the bottom borders and connections +/// +/// Each node is rendered in a box of fixed width (NODE_RENDER_WIDTH). +struct TreeRenderVisitor<'a, 'b> { + /// Write to this formatter + f: &'a mut Formatter<'b>, + /// Maximum total width of the rendered tree + maximum_render_width: usize, +} + +impl TreeRenderVisitor<'_, '_> { + // Unicode box-drawing characters for creating borders and connections. + const LTCORNER: &'static str = "┌"; // Left top corner + const RTCORNER: &'static str = "┐"; // Right top corner + const LDCORNER: &'static str = "└"; // Left bottom corner + const RDCORNER: &'static str = "┘"; // Right bottom corner + + const TMIDDLE: &'static str = "┬"; // Top T-junction (connects down) + const LMIDDLE: &'static str = "├"; // Left T-junction (connects right) + const DMIDDLE: &'static str = "┴"; // Bottom T-junction (connects up) + + const VERTICAL: &'static str = "│"; // Vertical line + const HORIZONTAL: &'static str = "─"; // Horizontal line + + // TODO: Make these variables configurable. + const NODE_RENDER_WIDTH: usize = 29; // Width of each node's box + const MAX_EXTRA_LINES: usize = 30; // Maximum number of extra info lines per node + + /// Main entry point for rendering an execution plan as a tree. + /// The rendering process happens in three stages for each level of the tree: + /// 1. Render top borders and connections + /// 2. Render node content and vertical connections + /// 3. Render bottom borders and connections + pub fn visit(&mut self, plan: &dyn ExecutionPlan) -> Result<(), fmt::Error> { + let root = RenderTree::create_tree(plan); + + for y in 0..root.height { + // Start by rendering the top layer. + self.render_top_layer(&root, y)?; + // Now we render the content of the boxes + self.render_box_content(&root, y)?; + // Render the bottom layer of each of the boxes + self.render_bottom_layer(&root, y)?; + } + + Ok(()) + } + + /// Renders the top layer of boxes at the given y-level of the tree. + /// This includes: + /// - Top corners (┌─┐) for nodes + /// - Horizontal connections between nodes + /// - Vertical connections to parent nodes + fn render_top_layer( + &mut self, + root: &RenderTree, + y: usize, + ) -> Result<(), fmt::Error> { + for x in 0..root.width { + if self.maximum_render_width > 0 + && x * Self::NODE_RENDER_WIDTH >= self.maximum_render_width + { + break; + } + + if root.has_node(x, y) { + write!(self.f, "{}", Self::LTCORNER)?; + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + if y == 0 { + // top level node: no node above this one + write!(self.f, "{}", Self::HORIZONTAL)?; + } else { + // render connection to node above this one + write!(self.f, "{}", Self::DMIDDLE)?; + } + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + write!(self.f, "{}", Self::RTCORNER)?; + } else { + let mut has_adjacent_nodes = false; + for i in 0..(root.width - x) { + has_adjacent_nodes = has_adjacent_nodes || root.has_node(x + i, y); + } + if !has_adjacent_nodes { + // There are no nodes to the right side of this position + // no need to fill the empty space + continue; + } + // there are nodes next to this, fill the space + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } + writeln!(self.f)?; + + Ok(()) + } + + /// Renders the content layer of boxes at the given y-level of the tree. + /// This includes: + /// - Node names and extra information + /// - Vertical borders (│) for boxes + /// - Vertical connections between nodes + fn render_box_content( + &mut self, + root: &RenderTree, + y: usize, + ) -> Result<(), fmt::Error> { + let mut extra_info: Vec<Vec<String>> = vec![vec![]; root.width]; + let mut extra_height = 0; + + for (x, extra_info_item) in extra_info.iter_mut().enumerate().take(root.width) { + if let Some(node) = root.get_node(x, y) { + Self::split_up_extra_info( + &node.extra_text, + extra_info_item, + Self::MAX_EXTRA_LINES, + ); + if extra_info_item.len() > extra_height { + extra_height = extra_info_item.len(); + } + } + } + + let halfway_point = extra_height.div_ceil(2); + + // Render the actual node. + for render_y in 0..=extra_height { + for (x, _) in root.nodes.iter().enumerate().take(root.width) { + if self.maximum_render_width > 0 + && x * Self::NODE_RENDER_WIDTH >= self.maximum_render_width + { + break; + } + + let mut has_adjacent_nodes = false; + for i in 0..(root.width - x) { + has_adjacent_nodes = has_adjacent_nodes || root.has_node(x + i, y); + } + + if let Some(node) = root.get_node(x, y) { + write!(self.f, "{}", Self::VERTICAL)?; + + // Figure out what to render. + let mut render_text = if render_y == 0 { + node.name.clone() + } else if render_y <= extra_info[x].len() { + extra_info[x][render_y - 1].clone() + } else { + String::new() + }; + + render_text = Self::adjust_text_for_rendering( + &render_text, + Self::NODE_RENDER_WIDTH - 2, + ); + write!(self.f, "{render_text}")?; + + if render_y == halfway_point && node.child_positions.len() > 1 { + write!(self.f, "{}", Self::LMIDDLE)?; + } else { + write!(self.f, "{}", Self::VERTICAL)?; + } + } else if render_y == halfway_point { + let has_child_to_the_right = + Self::should_render_whitespace(root, x, y); + if root.has_node(x, y + 1) { + // Node right below this one. + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + if has_child_to_the_right { + write!(self.f, "{}", Self::TMIDDLE)?; + // Have another child to the right, Keep rendering the line. + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + } else { + write!(self.f, "{}", Self::RTCORNER)?; + if has_adjacent_nodes { + // Only a child below this one: fill the reset with spaces. + write!( + self.f, + "{}", + " ".repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + } + } + } else if has_child_to_the_right { + // Child to the right, but no child right below this one: render a full + // line. + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH) + )?; + } else if has_adjacent_nodes { + // Empty spot: render spaces. + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } else if render_y >= halfway_point { + if root.has_node(x, y + 1) { + // Have a node below this empty spot: render a vertical line. + write!( + self.f, + "{}{}", + " ".repeat(Self::NODE_RENDER_WIDTH / 2), + Self::VERTICAL + )?; + if has_adjacent_nodes + || Self::should_render_whitespace(root, x, y) + { + write!( + self.f, + "{}", + " ".repeat(Self::NODE_RENDER_WIDTH / 2) + )?; + } + } else if has_adjacent_nodes + || Self::should_render_whitespace(root, x, y) + { + // Empty spot: render spaces. + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } else if has_adjacent_nodes { + // Empty spot: render spaces. + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } + writeln!(self.f)?; + } + + Ok(()) + } + + /// Renders the bottom layer of boxes at the given y-level of the tree. + /// This includes: + /// - Bottom corners (└─┘) for nodes + /// - Horizontal connections between nodes + /// - Vertical connections to child nodes + fn render_bottom_layer( + &mut self, + root: &RenderTree, + y: usize, + ) -> Result<(), fmt::Error> { + for x in 0..=root.width { + if self.maximum_render_width > 0 + && x * Self::NODE_RENDER_WIDTH >= self.maximum_render_width + { + break; + } + let mut has_adjacent_nodes = false; + for i in 0..(root.width - x) { + has_adjacent_nodes = has_adjacent_nodes || root.has_node(x + i, y); + } + if root.get_node(x, y).is_some() { + write!(self.f, "{}", Self::LDCORNER)?; + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + if root.has_node(x, y + 1) { + // node below this one: connect to that one + write!(self.f, "{}", Self::TMIDDLE)?; + } else { + // no node below this one: end the box + write!(self.f, "{}", Self::HORIZONTAL)?; + } + write!( + self.f, + "{}", + Self::HORIZONTAL.repeat(Self::NODE_RENDER_WIDTH / 2 - 1) + )?; + write!(self.f, "{}", Self::RDCORNER)?; + } else if root.has_node(x, y + 1) { + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH / 2))?; + write!(self.f, "{}", Self::VERTICAL)?; + if has_adjacent_nodes || Self::should_render_whitespace(root, x, y) { + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH / 2))?; + } + } else if has_adjacent_nodes || Self::should_render_whitespace(root, x, y) { + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; + } + } + writeln!(self.f)?; + + Ok(()) + } + + fn extra_info_separator() -> String { + "-".repeat(Self::NODE_RENDER_WIDTH - 9) + } + + fn remove_padding(s: &str) -> String { + s.trim().to_string() + } + + pub fn split_up_extra_info( + extra_info: &HashMap<String, String>, + result: &mut Vec<String>, + max_lines: usize, + ) { + if extra_info.is_empty() { + return; + } + + result.push(Self::extra_info_separator()); + + let mut requires_padding = false; + let mut was_inlined = false; + + // use BTreeMap for repeatable key order + let sorted_extra_info: BTreeMap<_, _> = extra_info.iter().collect(); + for (key, value) in sorted_extra_info { + let mut str = Self::remove_padding(value); + let mut is_inlined = false; + let available_width = Self::NODE_RENDER_WIDTH - 7; + let total_size = key.len() + str.len() + 2; + let is_multiline = str.contains('\n'); + + if str.is_empty() { + str = key.to_string(); + } else if !is_multiline && total_size < available_width { + str = format!("{key}: {str}"); + is_inlined = true; + } else { + str = format!("{key}:\n{str}"); + } + + if is_inlined && was_inlined { + requires_padding = false; + } + + if requires_padding { + result.push(String::new()); + } + + let mut splits: Vec<String> = str.split('\n').map(String::from).collect(); + if splits.len() > max_lines { + let mut truncated_splits = Vec::new(); + for split in splits.iter().take(max_lines / 2) { + truncated_splits.push(split.clone()); + } + truncated_splits.push("...".to_string()); + for split in splits.iter().skip(splits.len() - max_lines / 2) { + truncated_splits.push(split.clone()); + } + splits = truncated_splits; + } + for split in splits { + Self::split_string_buffer(&split, result); + } + if result.len() > max_lines { + result.truncate(max_lines); + result.push("...".to_string()); + } + + requires_padding = true; + was_inlined = is_inlined; + } + } + + /// Adjusts text to fit within the specified width by: + /// 1. Truncating with ellipsis if too long + /// 2. Center-aligning within the available space if shorter + fn adjust_text_for_rendering(source: &str, max_render_width: usize) -> String { + let render_width = source.chars().count(); + if render_width > max_render_width { + let truncated = &source[..max_render_width - 3]; + format!("{truncated}...") + } else { + let total_spaces = max_render_width - render_width; + let half_spaces = total_spaces / 2; + let extra_left_space = if total_spaces.is_multiple_of(2) { 0 } else { 1 }; + format!( + "{}{}{}", + " ".repeat(half_spaces + extra_left_space), + source, + " ".repeat(half_spaces) + ) + } + } + + /// Determines if whitespace should be rendered at a given position. + /// This is important for: + /// 1. Maintaining proper spacing between sibling nodes + /// 2. Ensuring correct alignment of connections between parents and children + /// 3. Preserving the tree structure's visual clarity + fn should_render_whitespace(root: &RenderTree, x: usize, y: usize) -> bool { + let mut found_children = 0; + + for i in (0..=x).rev() { + let node = root.get_node(i, y); + if root.has_node(i, y + 1) { + found_children += 1; + } + if let Some(node) = node { + if node.child_positions.len() > 1 + && found_children < node.child_positions.len() + { + return true; + } + + return false; + } + } + + false + } + + fn split_string_buffer(source: &str, result: &mut Vec<String>) { + let mut character_pos = 0; + let mut start_pos = 0; + let mut render_width = 0; + let mut last_possible_split = 0; + + let chars: Vec<char> = source.chars().collect(); + + while character_pos < chars.len() { + // Treating each char as width 1 for simplification + let char_width = 1; + + // Does the next character make us exceed the line length? + if render_width + char_width > Self::NODE_RENDER_WIDTH - 2 { + if start_pos + 8 > last_possible_split { + // The last character we can split on is one of the first 8 characters of the line + // to not create very small lines we instead split on the current character + last_possible_split = character_pos; + } + + result.push(source[start_pos..last_possible_split].to_string()); + render_width = character_pos - last_possible_split; + start_pos = last_possible_split; + character_pos = last_possible_split; + } + + // check if we can split on this character + if Self::can_split_on_this_char(chars[character_pos]) { + last_possible_split = character_pos; + } + + character_pos += 1; + render_width += char_width; + } + + if source.len() > start_pos { + // append the remainder of the input + result.push(source[start_pos..].to_string()); + } + } + + fn can_split_on_this_char(c: char) -> bool { + (!c.is_ascii_digit() && !c.is_ascii_uppercase() && !c.is_ascii_lowercase()) + && c != '_' + } +} + +/// Trait for types which could have additional details when formatted in `Verbose` mode +pub trait DisplayAs { + /// Format according to `DisplayFormatType`, used when verbose representation looks + /// different from the default one + /// + /// Should not include a newline + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result; +} + +/// A new type wrapper to display `T` implementing`DisplayAs` using the `Default` mode +pub struct DefaultDisplay<T>(pub T); + +impl<T: DisplayAs> fmt::Display for DefaultDisplay<T> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } +} + +/// A new type wrapper to display `T` implementing `DisplayAs` using the `Verbose` mode +pub struct VerboseDisplay<T>(pub T); + +impl<T: DisplayAs> fmt::Display for VerboseDisplay<T> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Verbose, f) + } +} + +/// A wrapper to customize partitioned file display +#[derive(Debug)] +pub struct ProjectSchemaDisplay<'a>(pub &'a SchemaRef); + +impl fmt::Display for ProjectSchemaDisplay<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let parts: Vec<_> = self + .0 + .fields() + .iter() + .map(|x| x.name().to_owned()) + .collect::<Vec<String>>(); + write!(f, "[{}]", parts.join(", ")) + } +} + +pub fn display_orderings(f: &mut Formatter, orderings: &[LexOrdering]) -> fmt::Result { + if !orderings.is_empty() { + let start = if orderings.len() == 1 { + ", output_ordering=" + } else { + ", output_orderings=[" + }; + write!(f, "{start}")?; + for (idx, ordering) in orderings.iter().enumerate() { + match idx { + 0 => write!(f, "[{ordering}]")?, + _ => write!(f, ", [{ordering}]")?, + } + } + let end = if orderings.len() == 1 { "" } else { "]" }; + write!(f, "{end}")?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::fmt::Write; + use std::sync::Arc; + + use datafusion_common::{ + Result, Statistics, internal_datafusion_err, tree_node::TreeNodeRecursion, + }; + use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + use datafusion_physical_expr::PhysicalExpr; + + use crate::statistics::StatisticsArgs; + use crate::{ + ChildrenPropertiesMode, DisplayAs, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, + }; + + use super::DisplayableExecutionPlan; + + #[derive(Debug, Clone, Copy)] + enum TestStatsExecPlan { + Panic, + Error, + Ok, + } + + impl DisplayAs for TestStatsExecPlan { + fn fmt_as( + &self, + _t: crate::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "TestStatsExecPlan") + } + } + + impl ExecutionPlan for TestStatsExecPlan { + fn name(&self) -> &'static str { + "TestStatsExecPlan" + } + + fn properties(&self) -> &Arc<PlanProperties> { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _: usize, + _: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + todo!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if args.partition().is_some() { + return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))); + } + match self { + Self::Panic => panic!("expected panic"), + Self::Error => Err(internal_datafusion_err!("expected error")), + Self::Ok => Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))), + } + } + } + + fn test_stats_display(exec: TestStatsExecPlan, show_stats: bool) { + let display = + DisplayableExecutionPlan::new(&exec).set_show_statistics(show_stats); + + let mut buf = String::new(); + write!(&mut buf, "{}", display.one_line()).unwrap(); + let buf = buf.trim(); + assert_eq!(buf, "TestStatsExecPlan"); + } + + #[test] + fn test_display_when_stats_panic_with_no_show_stats() { + test_stats_display(TestStatsExecPlan::Panic, false); + } + + #[test] + fn test_display_when_stats_error_with_no_show_stats() { + test_stats_display(TestStatsExecPlan::Error, false); + } + + #[test] + fn test_display_when_stats_ok_with_no_show_stats() { + test_stats_display(TestStatsExecPlan::Ok, false); + } + + #[test] + #[should_panic(expected = "expected panic")] + fn test_display_when_stats_panic_with_show_stats() { + test_stats_display(TestStatsExecPlan::Panic, true); + } + + #[test] + #[should_panic(expected = "Error")] // fmt::Error + fn test_display_when_stats_error_with_show_stats() { + test_stats_display(TestStatsExecPlan::Error, true); + } + + #[test] + fn test_display_when_stats_ok_with_show_stats() { + test_stats_display(TestStatsExecPlan::Ok, false); + } + + mod pgjson { + use std::sync::Arc; + use std::time::Duration; + + use arrow::datatypes::{DataType, Field, Schema}; + use insta::assert_snapshot; + + use super::super::DisplayableExecutionPlan; + use crate::empty::EmptyExec; + use crate::filter::FilterExec; + use crate::projection::ProjectionExec; + use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; + use datafusion_physical_expr::expressions::{binary, col, lit}; + use datafusion_physical_expr::{Partitioning, PhysicalExpr}; + + fn sample_plan() -> Arc<dyn ExecutionPlan> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let predicate = binary( + col("a", &schema).unwrap(), + datafusion_expr::Operator::Gt, + lit(5i32), + &schema, + ) + .unwrap(); + let filter = Arc::new(FilterExec::try_new(predicate, empty).unwrap()); + let proj_expr: Vec<(Arc<dyn PhysicalExpr>, String)> = + vec![(col("a", &schema).unwrap(), "a".to_string())]; + let _ = Partitioning::UnknownPartitioning(1); + Arc::new(ProjectionExec::try_new(proj_expr, filter).unwrap()) + } + + #[test] + fn pgjson_renders_plan_without_metrics() { + let plan = sample_plan(); + let out = DisplayableExecutionPlan::new(plan.as_ref()) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + // Root is an array with one {"Plan": ...} entry. + let root = value + .as_array() + .expect("root array") + .first() + .expect("root entry") + .get("Plan") + .expect("plan object"); + assert_eq!(root["Node Type"].as_str(), Some("ProjectionExec")); + assert!(root.get("Actual Rows").is_none()); + assert!(root.get("Extras").is_none()); + let plans = root["Plans"].as_array().expect("Plans array"); + assert_eq!(plans.len(), 1); + assert_eq!(plans[0]["Node Type"].as_str(), Some("FilterExec")); + } + + #[test] + fn pgjson_emits_pg_canonical_metric_keys() { + use crate::metrics::{Count, Metric, MetricValue, MetricsSet, Time}; + use crate::{DisplayFormatType, ExecutionPlan, PlanProperties}; + use datafusion_common::Result; + use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + + // Wrap `sample_plan()` with an adapter node that exposes a + // hand-crafted `MetricsSet` so we can assert the PG key mapping + // without running anything. + #[derive(Debug)] + struct WithMetrics { + inner: Arc<dyn ExecutionPlan>, + metrics: MetricsSet, + } + impl crate::DisplayAs for WithMetrics { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "WithMetrics") + } + } + impl ExecutionPlan for WithMetrics { + fn name(&self) -> &'static str { + "WithMetrics" + } + fn properties(&self) -> &Arc<PlanProperties> { + self.inner.properties() + } + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.inner] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut( + &Arc<dyn PhysicalExpr>, + ) -> Result< + datafusion_common::tree_node::TreeNodeRecursion, + >, + ) -> Result<datafusion_common::tree_node::TreeNodeRecursion> + { + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( + &self, + _: usize, + _: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone()) + } + } + + let mut metrics = MetricsSet::new(); + let rows = Count::new(); + rows.add(42); + metrics.push(Arc::new(Metric::new(MetricValue::OutputRows(rows), None))); + let elapsed = Time::new(); + elapsed.add_duration(Duration::from_millis(5)); + metrics.push(Arc::new(Metric::new( + MetricValue::ElapsedCompute(elapsed), + None, + ))); + let batches = Count::new(); + batches.add(7); + metrics.push(Arc::new(Metric::new( + MetricValue::OutputBatches(batches), + None, + ))); + + let plan: Arc<dyn ExecutionPlan> = Arc::new(WithMetrics { + inner: sample_plan(), + metrics, + }); + + let out = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + let root = value[0].get("Plan").expect("plan"); + assert_eq!(root["Actual Rows"].as_u64(), Some(42)); + assert_eq!(root["Actual Total Time"].as_f64(), Some(5.0)); + assert_eq!(root["Extras"]["output_batches"].as_u64(), Some(7)); + + let metric_names = vec!["output_rows".to_string()]; + for rendered in [ + DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .indent(false) + .to_string(), + DisplayableExecutionPlan::with_full_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .indent(false) + .to_string(), + DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .graphviz() + .to_string(), + DisplayableExecutionPlan::with_full_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .graphviz() + .to_string(), + ] { + assert!(rendered.contains("output_rows")); + assert!(!rendered.contains("elapsed_compute")); + assert!(!rendered.contains("output_batches")); + } + + let out = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(metric_names) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + let root = value[0].get("Plan").expect("plan"); + assert_eq!(root["Actual Rows"].as_u64(), Some(42)); + assert!(root.get("Actual Total Time").is_none()); + assert!(root.get("Extras").is_none()); + } + + #[test] + fn pgjson_includes_summary_when_set() { + let plan = sample_plan(); + let out = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_summary(Some(42), Some(Duration::from_millis(7))) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + let entry = &value.as_array().unwrap()[0]; + assert_eq!(entry["Total Rows"].as_u64(), Some(42)); + assert!(entry["Duration"].is_string()); + } + + #[test] + fn pgjson_snapshot_of_sample_plan() { + let plan = sample_plan(); + let out = DisplayableExecutionPlan::new(plan.as_ref()) + .pgjson(false) + .to_string(); + // This snapshot assumes `serde_json` is built with the + // `preserve_order` feature (enabled via this crate's dev-deps). + assert_snapshot!(out, @r#" + [ + { + "Plan": { + "Node Type": "ProjectionExec", + "Details": "ProjectionExec: expr=[a@0 as a]", + "Plans": [ + { + "Node Type": "FilterExec", + "Details": "FilterExec: a@0 > 5", + "Plans": [ + { + "Node Type": "EmptyExec", + "Details": "EmptyExec", + "Plans": [] + } + ] + } + ] + } + } + ] + "#); + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/distribution_requirements.rs b/native/vendor/datafusion-physical-plan/src/distribution_requirements.rs new file mode 100644 index 0000000..6405b1f --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/distribution_requirements.rs
@@ -0,0 +1,359 @@ +// 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. + +//! Input distribution requirements for physical execution plans. + +use datafusion_common::{Result, internal_err}; +use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction}; + +use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; + +/// Distribution requirements for an [`ExecutionPlan`]'s inputs. +/// +/// [`InputDistributionRequirements`] describes what distribution an operator +/// requires from each child. +/// +/// - [`Self::new`] describes independent per-child requirements. +/// - [`Self::co_partitioned`] additionally requires child partitions with the +/// same index to cover compatible key ranges. +/// +/// For a single-input aggregate: +/// +/// ```text +/// AggregateExec +/// child 0 requirement: KeyPartitioned(group_exprs) +/// ``` +/// +/// each input partition can aggregate its own key domain independently. +/// +/// For a partitioned join: +/// +/// ```text +/// HashJoinExec +/// child 0 requirement: KeyPartitioned(left_keys) +/// child 1 requirement: KeyPartitioned(right_keys) +/// +/// partition 0: join(left partition 0, right partition 0) +/// partition 1: join(left partition 1, right partition 1) +/// partition 2: join(left partition 2, right partition 2) +/// ``` +/// +/// each child must satisfy its own key requirement. In addition, matching +/// partition indexes must be safe to process together. +#[non_exhaustive] +#[derive(Debug, Clone)] +pub struct InputDistributionRequirements { + /// Per-child distribution requirements, indexed by child position. + children: Vec<ChildDistributionRequirement>, + /// Child indexes that must also have compatible partition layouts. + co_partitioned: Option<Vec<usize>>, +} + +/// Options for checking child distribution satisfaction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ChildSatisfactionOptions { + allow_subset: bool, +} + +impl ChildSatisfactionOptions { + /// Create default satisfaction options. + pub fn new() -> Self { + Self::default() + } + + /// Allow a child partitioning whose key expressions are a subset of the + /// required key expressions to satisfy the requirement. + pub fn with_allow_subset(mut self, allow_subset: bool) -> Self { + self.allow_subset = allow_subset; + self + } + + /// Whether subset satisfaction is enabled. + pub fn allow_subset(&self) -> bool { + self.allow_subset + } +} + +impl InputDistributionRequirements { + /// Create independent per-child requirements. + pub fn new(per_child: Vec<Distribution>) -> Self { + let children = per_child + .into_iter() + .map(|distribution| ChildDistributionRequirement { distribution }) + .collect(); + + Self { + children, + co_partitioned: None, + } + } + + /// Create a requirement that all children are co-partitioned. + /// + /// Each child must satisfy its own [`Distribution`]. Matching partition + /// indexes are processed together: + /// + /// ```text + /// left: Range(left.a ASC, split_points=[10, 20]) + /// right: Range(right.x ASC, split_points=[10, 20]) + /// + /// partition 0 from both sides contains keys before 10 + /// partition 1 from both sides contains keys in [10, 20) + /// partition 2 from both sides contains keys at/after 20 + /// ``` + /// + /// If the split points differ, partition `i` from one side no longer covers + /// the same key range as partition `i` from the other side. + pub fn co_partitioned(per_child: Vec<Distribution>) -> Self { + debug_assert!( + per_child.len() >= 2, + "co-partitioned distribution requirements need at least two children" + ); + let co_partitioned = (0..per_child.len()).collect(); + let mut result = Self::new(per_child); + result.co_partitioned = Some(co_partitioned); + result + } + + /// Return the per-child distribution requirements. + pub fn per_child_distributions( + &self, + ) -> impl ExactSizeIterator<Item = &Distribution> + '_ { + self.children.iter().map(|child| &child.distribution) + } + + /// Return the distribution requirement for a child. + pub fn child_distribution(&self, child_idx: usize) -> Option<&Distribution> { + self.children + .get(child_idx) + .map(|child| &child.distribution) + } + + /// Return the per-child distribution requirements. + /// + /// WARNING: This intentionally drops any grouped relationship. + pub fn into_per_child(self) -> Vec<Distribution> { + self.children + .into_iter() + .map(|child| child.distribution) + .collect() + } + + /// Returns how a child satisfies its distribution requirement. + /// + /// This preserves the requirement set's satisfaction policy. + pub fn child_satisfaction( + &self, + child_idx: usize, + child: &dyn ExecutionPlan, + options: ChildSatisfactionOptions, + ) -> Result<PartitioningSatisfaction> { + let Some(requirement) = self.children.get(child_idx) else { + return internal_err!( + "missing distribution requirement for child {child_idx}" + ); + }; + + Ok(child.output_partitioning().satisfaction( + &requirement.distribution, + child.equivalence_properties(), + options.allow_subset(), + )) + } + + /// Return child indexes whose co-partitioning requirements are + /// unsatisfied by the provided candidate children. + /// + /// Independent per-child requirements are intentionally ignored here, use + /// [`Self::child_satisfaction`] for those checks. An empty result means all + /// co-partitioning requirements are satisfied. + #[doc(hidden)] + pub fn unsatisfied_co_partitioned_children( + &self, + plan_name: &str, + children: &[&dyn ExecutionPlan], + ) -> Result<Vec<usize>> { + self.validate_shape(plan_name, children.len())?; + + let Some(co_partitioned) = &self.co_partitioned else { + return Ok(vec![]); + }; + if self.co_partitioning_satisfied(co_partitioned, children) { + return Ok(vec![]); + } + + Ok(co_partitioned.clone()) + } + + /// Validate the requirements against a plan's children. + pub(crate) fn check_invariants<P: ExecutionPlan + ?Sized>( + &self, + plan: &P, + check: InvariantLevel, + ) -> Result<()> { + let children = plan.children(); + self.validate_shape(plan.name(), children.len())?; + + let children = children + .into_iter() + .map(|child| child.as_ref()) + .collect::<Vec<_>>(); + if matches!(check, InvariantLevel::Executable) + && let Some(co_partitioned) = &self.co_partitioned + && !self.co_partitioning_satisfied(co_partitioned, &children) + { + return internal_err!( + "{} requires children {:?} to be co-partitioned", + plan.name(), + co_partitioned + ); + } + + Ok(()) + } + + fn validate_shape(&self, plan_name: &str, children_len: usize) -> Result<()> { + if self.children.len() != children_len { + return internal_err!( + "{plan_name}::input_distribution_requirements returned incorrect child count: {} != {}", + self.children.len(), + children_len + ); + } + + if let Some(co_partitioned) = &self.co_partitioned { + if co_partitioned.len() < 2 { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: at least two children are required" + ); + } + let mut seen = vec![false; self.children.len()]; + for &child in co_partitioned { + validate_child_index(plan_name, child, self.children.len(), &mut seen)?; + if matches!( + self.children[child].distribution, + Distribution::UnspecifiedDistribution + ) { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: child {child} has unspecified distribution" + ); + } + } + } + + Ok(()) + } + + fn co_partitioning_satisfied( + &self, + co_partitioned: &[usize], + children: &[&dyn ExecutionPlan], + ) -> bool { + let first_idx = co_partitioned[0]; + let first_requirement = &self.children[first_idx]; + let first = children[first_idx]; + let first_partitioning = first.output_partitioning(); + + if !first_partitioning + .satisfaction( + &first_requirement.distribution, + first.equivalence_properties(), + false, + ) + .is_satisfied() + { + return false; + } + + for &child_idx in co_partitioned.iter().skip(1) { + let requirement = &self.children[child_idx]; + let child = children[child_idx]; + if !child + .output_partitioning() + .satisfaction( + &requirement.distribution, + child.equivalence_properties(), + false, + ) + .is_satisfied() + || !compatible_co_partitioning_layout( + first_partitioning, + child.output_partitioning(), + ) + { + return false; + } + } + + true + } +} + +/// A distribution requirement for a single child. +#[derive(Debug, Clone)] +struct ChildDistributionRequirement { + distribution: Distribution, +} + +fn validate_child_index( + plan_name: &str, + child_idx: usize, + child_count: usize, + seen: &mut [bool], +) -> Result<()> { + if child_idx >= child_count { + return internal_err!( + "{plan_name} has invalid distribution requirement: child index {child_idx} out of bounds" + ); + } + if seen[child_idx] { + return internal_err!( + "{plan_name} has invalid distribution requirement: child {child_idx} appears more than once" + ); + } + seen[child_idx] = true; + Ok(()) +} + +fn compatible_co_partitioning_layout( + first_partitioning: &Partitioning, + other_partitioning: &Partitioning, +) -> bool { + if first_partitioning.partition_count() == 1 + && other_partitioning.partition_count() == 1 + { + return true; + } + + if first_partitioning.partition_count() != other_partitioning.partition_count() { + return false; + } + + match (first_partitioning, other_partitioning) { + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, + (Partitioning::Range(left), Partitioning::Range(right)) => { + left.split_points() == right.split_points() + && left.ordering().len() == right.ordering().len() + && left + .ordering() + .iter() + .zip(right.ordering()) + .all(|(left, right)| left.options == right.options) + } + _ => false, + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/empty.rs b/native/vendor/datafusion-physical-plan/src/empty.rs new file mode 100644 index 0000000..dd08ff3 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/empty.rs
@@ -0,0 +1,313 @@ +// 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. + +//! EmptyRelation with produce_one_row=false execution plan + +use std::sync::Arc; + +use crate::memory::MemoryStream; +use crate::{ + ChildrenPropertiesMode, DisplayAs, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, +}; +use crate::{ + DisplayFormatType, ExecutionPlan, Partitioning, + execution_plan::{Boundedness, EmissionType}, +}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ColumnStatistics, Result, ScalarValue, assert_or_internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; + +use crate::execution_plan::SchedulingType; +use crate::statistics::StatisticsArgs; +use log::trace; + +/// Execution plan for empty relation with produce_one_row=false +#[derive(Debug, Clone)] +pub struct EmptyExec { + /// The schema for the produced row + schema: SchemaRef, + /// Number of partitions + partitions: usize, + cache: Arc<PlanProperties>, +} + +impl EmptyExec { + /// Create a new EmptyExec + pub fn new(schema: SchemaRef) -> Self { + let cache = Self::compute_properties(Arc::clone(&schema), 1); + EmptyExec { + schema, + partitions: 1, + cache: Arc::new(cache), + } + } + + /// Create a new EmptyExec with specified partition number + pub fn with_partitions(mut self, partitions: usize) -> Self { + self.partitions = partitions; + // Changing partitions may invalidate output partitioning, so update it: + let output_partitioning = Self::output_partitioning_helper(self.partitions); + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + self + } + + fn data(&self) -> Result<Vec<RecordBatch>> { + Ok(vec![]) + } + + fn output_partitioning_helper(n_partitions: usize) -> Partitioning { + Partitioning::UnknownPartitioning(n_partitions) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef, n_partitions: usize) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Self::output_partitioning_helper(n_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + ) + .with_scheduling_type(SchedulingType::Cooperative) + } +} + +impl DisplayAs for EmptyExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "EmptyExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for EmptyExec { + fn name(&self) -> &'static str { + "EmptyExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start EmptyExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + assert_or_internal_err!( + partition < self.partitions, + "EmptyExec invalid partition {} (expected less than {})", + partition, + self.partitions + ); + + Ok(Box::pin(MemoryStream::try_new( + self.data()?, + Arc::clone(&self.schema), + None, + )?)) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if let Some(partition) = args.partition() { + assert_or_internal_err!( + partition < self.partitions, + "EmptyExec invalid partition {} (expected less than {})", + partition, + self.partitions + ); + } + + // Build explicit stats: exact zero rows and bytes, with explicit known column stats + let mut stats = Statistics::default() + .with_num_rows(Precision::Exact(0)) + .with_total_byte_size(Precision::Exact(0)); + + // Add explicit column stats for each field in schema + for _ in self.schema.fields() { + stats = stats.add_column_statistics(ColumnStatistics { + null_count: Precision::Exact(0), + distinct_count: Precision::Exact(0), + min_value: Precision::<ScalarValue>::Absent, + max_value: Precision::<ScalarValue>::Absent, + sum_value: Precision::<ScalarValue>::Absent, + byte_size: Precision::Exact(0), + }); + } + + Ok(Arc::new(stats)) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(schema), + partitions: self + .properties() + .output_partitioning() + .partition_count() as u32, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl EmptyExec { + /// Reconstruct an [`EmptyExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let empty = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Empty, + "EmptyExec", + ); + let schema = empty.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "EmptyExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?); + // A zero (absent) partition count comes from a plan encoded before the + // field existed, which always meant a single partition. + let partitions = empty.partitions.max(1) as usize; + Ok(Arc::new(EmptyExec::new(schema).with_partitions(partitions))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common; + use crate::execution_plan::replace_children_if_necessary; + use crate::test; + + #[tokio::test] + async fn empty() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = test::aggr_test_schema(); + + let empty = EmptyExec::new(Arc::clone(&schema)); + assert_eq!(empty.schema(), schema); + + // We should have no results + let iter = empty.execute(0, task_ctx)?; + let batches = common::collect(iter).await?; + assert!(batches.is_empty()); + + Ok(()) + } + + #[test] + fn with_new_children() -> Result<()> { + let schema = test::aggr_test_schema(); + let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let empty2 = replace_children_if_necessary( + Arc::clone(&empty) as Arc<dyn ExecutionPlan>, + vec![], + )?; + assert_eq!(empty.schema(), empty2.schema()); + + let too_many_kids = vec![empty2]; + assert!( + replace_children_if_necessary(empty, too_many_kids).is_err(), + "expected error when providing list of kids" + ); + Ok(()) + } + + #[tokio::test] + async fn invalid_execute() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = test::aggr_test_schema(); + let empty = EmptyExec::new(schema); + + // ask for the wrong partition + assert!(empty.execute(1, Arc::clone(&task_ctx)).is_err()); + assert!(empty.execute(20, task_ctx).is_err()); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/execution_plan.rs b/native/vendor/datafusion-physical-plan/src/execution_plan.rs new file mode 100644 index 0000000..a4d081b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/execution_plan.rs
@@ -0,0 +1,3077 @@ +// 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. + +pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +use crate::distribution_requirements::InputDistributionRequirements; +use crate::filter_pushdown::{ + ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +pub use crate::metrics::Metric; +pub use crate::ordering::InputOrderMode; +use crate::sort_pushdown::SortOrderPushdownResult; +pub use crate::stream::EmptyRecordBatchStream; + +use arrow_schema::Schema; +pub use datafusion_common::hash_utils; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; +pub use datafusion_common::utils::project_schema; +pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; +pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +pub use datafusion_expr::{Accumulator, ColumnarValue}; +use datafusion_physical_expr::projection::ProjectionExpr; +pub use datafusion_physical_expr::window::WindowExpr; +pub use datafusion_physical_expr::{ + Distribution, Partitioning, PhysicalExpr, expressions, +}; + +use std::any::Any; +use std::collections::HashSet; +use std::fmt::Debug; +use std::sync::{Arc, LazyLock}; + +use crate::coalesce_partitions::CoalescePartitionsExec; +use crate::display::DisplayableExecutionPlan; +use crate::metrics::MetricsSet; +use crate::projection::ProjectionExec; +use crate::repartition::RepartitionExec; +use crate::sorts::sort_preserving_merge::SortPreservingMergeExec; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; + +use arrow::array::{Array, RecordBatch}; +use arrow::datatypes::SchemaRef; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{ + Constraints, DataFusionError, Result, assert_eq_or_internal_err, + assert_or_internal_err, exec_err, +}; +use datafusion_common_runtime::JoinSet; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr_common::sort_expr::{ + LexOrdering, OrderingRequirements, PhysicalSortExpr, +}; + +use futures::stream::{StreamExt, TryStreamExt}; + +/// Represent nodes in the DataFusion Physical Plan. +/// +/// Calling [`execute`] produces an `async` [`SendableRecordBatchStream`] of +/// [`RecordBatch`] that incrementally computes a partition of the +/// `ExecutionPlan`'s output from its input. See [`Partitioning`] for more +/// details on partitioning. +/// +/// Methods such as [`Self::schema`] and [`Self::properties`] communicate +/// properties of the output to the DataFusion optimizer, and methods such as +/// [`required_input_distribution`] and [`required_input_ordering`] express +/// requirements of the `ExecutionPlan` from its input. +/// +/// [`ExecutionPlan`] can be displayed in a simplified form using the +/// return value from [`displayable`] in addition to the (normally +/// quite verbose) `Debug` output. +/// +/// [`execute`]: ExecutionPlan::execute +/// [`required_input_distribution`]: ExecutionPlan::required_input_distribution +/// [`required_input_ordering`]: ExecutionPlan::required_input_ordering +/// +/// # Examples +/// +/// See [`datafusion-examples`] for examples, including +/// [`memory_pool_execution_plan.rs`] which shows how to implement a custom +/// `ExecutionPlan` with memory tracking and spilling support. +/// +/// [`datafusion-examples`]: https://github.com/apache/datafusion/tree/main/datafusion-examples +/// [`memory_pool_execution_plan.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { + /// Short name for the ExecutionPlan, such as 'DataSourceExec'. + /// + /// Implementation note: this method can just proxy to + /// [`static_name`](ExecutionPlan::static_name) if no special action is + /// needed. It doesn't provide a default implementation like that because + /// this method doesn't require the `Sized` constrain to allow a wilder + /// range of use cases. + fn name(&self) -> &str; + + /// Short name for the ExecutionPlan, such as 'DataSourceExec'. + /// Like [`name`](ExecutionPlan::name) but can be called without an instance. + fn static_name() -> &'static str + where + Self: Sized, + { + let full_name = std::any::type_name::<Self>(); + let maybe_start_idx = full_name.rfind(':'); + match maybe_start_idx { + Some(start_idx) => &full_name[start_idx + 1..], + None => "UNKNOWN", + } + } + + /// Returns the plan that provides this plan's public + /// [`ExecutionPlan`] downcast identity. + /// + /// This hook is for wrapper nodes that delegate their public downcast + /// identity to another plan while adding cross-cutting behavior such as + /// instrumentation. The default implementation returns `None`, meaning this + /// plan's concrete type is used for type introspection. + /// + /// Most `ExecutionPlan` implementations should use the default `None`; + /// override this only for wrapper plans that intentionally delegate their + /// public downcast identity to another plan. + /// + /// The `is` and `downcast_ref` helpers follow the returned delegate instead + /// of checking the current concrete type, making intermediate delegating + /// wrappers invisible to normal downcast-based inspection. + /// + /// Implementations that opt in should return the delegate plan, not `self`. + /// + /// This is independent from [`Self::children`] and should not be used for + /// plan traversal or optimizer rewrites. + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + None + } + + /// Get the schema for this execution plan + fn schema(&self) -> SchemaRef { + Arc::clone(self.properties().schema()) + } + + /// Return properties of the output of the `ExecutionPlan`, such as output + /// ordering(s), partitioning information etc. + /// + /// This information is available via methods on [`ExecutionPlanProperties`] + /// trait, which is implemented for all `ExecutionPlan`s. + fn properties(&self) -> &Arc<PlanProperties>; + + /// Returns an error if this individual node does not conform to its invariants. + /// These invariants are typically only checked in debug mode. + /// + /// A default set of invariants is provided in the [check_default_invariants] function. + /// The default implementation of `check_invariants` calls this function. + /// Extension nodes can provide their own invariants. + fn check_invariants(&self, check: InvariantLevel) -> Result<()> { + check_default_invariants(self, check) + } + + /// Returns the dynamic expressions produced by this plan node. + /// + /// A dynamic expression is produced when this node updates or completes its + /// runtime state during execution. Expressions that this node only consumes + /// must not be returned. This method is shallow and does not include dynamic + /// expressions produced by child plans. + /// + /// Each returned expression must have a [`PhysicalExpr::expression_id`] + /// since all dynamic expressions such as [`DynamicFilterPhysicalExpr`] + /// have an expression id. + /// + /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr + fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> { + Vec::new() + } + + /// Specifies simple per-child input distribution requirements. + /// + /// Deprecated: override [`Self::input_distribution_requirements`] instead. + /// + /// By default, each child has [`Distribution::UnspecifiedDistribution`]. + #[deprecated(since = "55.0.0", note = "Use input_distribution_requirements")] + fn required_input_distribution(&self) -> Vec<Distribution> { + vec![Distribution::UnspecifiedDistribution; self.children().len()] + } + + /// Specifies the input distribution requirements for this plan. + /// + /// The default implementation wraps [`Self::required_input_distribution`]. + /// Override this method for richer requirements, such as allowing alternate + /// satisfaction policies or requiring multiple children to be co-partitioned. + /// See [`InputDistributionRequirements`] for details. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + #[expect( + deprecated, + reason = "compatibility shim for external ExecutionPlan implementations" + )] + InputDistributionRequirements::new(self.required_input_distribution()) + } + + /// Specifies the ordering required for all of the children of this + /// `ExecutionPlan`. + /// + /// For each child, it's the local ordering requirement within + /// each partition rather than the global ordering + /// + /// NOTE that checking `!is_empty()` does **not** check for a + /// required input ordering. Instead, the correct check is that at + /// least one entry must be `Some` + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![None; self.children().len()] + } + + /// Returns `false` if this `ExecutionPlan`'s implementation may reorder + /// rows within or between partitions. + /// + /// For example, Projection, Filter, and Limit maintain the order + /// of inputs -- they may transform values (Projection) or not + /// produce the same number of rows that went in (Filter and + /// Limit), but the rows that are produced go in the same way. + /// + /// DataFusion uses this metadata to apply certain optimizations + /// such as automatically repartitioning correctly. + /// + /// The default implementation returns `false` + /// + /// WARNING: if you override this default, you *MUST* ensure that + /// the `ExecutionPlan`'s maintains the ordering invariant or else + /// DataFusion may produce incorrect results. + fn maintains_input_order(&self) -> Vec<bool> { + vec![false; self.children().len()] + } + + /// Specifies whether the `ExecutionPlan` benefits from increased + /// parallelization at its input for each child. + /// + /// If returns `true`, the `ExecutionPlan` would benefit from partitioning + /// its corresponding child (and thus from more parallelism). For + /// `ExecutionPlan` that do very little work the overhead of extra + /// parallelism may outweigh any benefits + /// + /// The default implementation returns `true` unless this `ExecutionPlan` + /// has signalled it requires a single child input partition. + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + // By default try to maximize parallelism with more CPUs if + // possible + self.input_distribution_requirements() + .per_child_distributions() + .map(|dist| !matches!(dist, Distribution::SinglePartition)) + .collect() + } + + /// Get a list of children `ExecutionPlan`s that act as inputs to this plan. + /// The returned list will be empty for leaf nodes such as scans, will contain + /// a single value for unary nodes, or two values for binary nodes (such as + /// joins). + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>>; + + /// Returns a clone of the existing plan with the children replaced, + /// skipping recomputation of plan properties when the options indicate + /// the new children's properties are unchanged. + /// + /// Callers should typically call [`replace_children_if_necessary`] and + /// not invoke this method directly. + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + #[expect(deprecated)] + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.with_new_children_and_same_properties(children) + } + ChildrenPropertiesMode::Recompute => self.with_new_children(children), + } + } + + /// Apply a closure `f` to each root expression that this node owns and uses + /// during execution, either by evaluating it or updating it dynamically. + /// + /// An expression must not be visited solely because it describes an input or + /// output property, such as cached ordering, partitioning, or equivalence + /// metadata. However, these may be traversed indirectly. For example, + /// `RepartitionExec` visits the partitioning expressions it evaluates and + /// `SortExec` visits the sort expressions it evaluates to order rows. + /// + /// This method is shallow: it must not visit expression children or expressions + /// owned by child execution plans. + /// + /// Similarly to other [`TreeNode`] APIs, the closure can return + /// [`TreeNodeRecursion::Stop`] to stop iteration, otherwise iteration + /// should continue. Note that [`TreeNodeRecursion::Continue`] and + /// [`TreeNodeRecursion::Jump`] are equivalent because this method is not + /// recursive. + /// + /// + /// # Example Usage + /// ``` + /// # use std::sync::Arc; + /// # use datafusion_physical_plan::ExecutionPlan; + /// # use datafusion_common::tree_node::TreeNodeRecursion; + /// # fn example(plan: Arc<dyn ExecutionPlan>) -> datafusion_common::Result<()> { + /// // Count the number of expressions + /// let mut count = 0; + /// plan.apply_expressions(&mut |_expr| { + /// count += 1; + /// Ok(TreeNodeRecursion::Continue) + /// })?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Implementation Examples + /// + /// ## Node with expressions (e.g., FilterExec, ProjectionExec) + /// + /// Use [`apply_expression_roots`] to implement this method. It abstracts away the + /// [`TreeNodeRecursion`] iteration from implementors. + /// ```ignore + /// fn apply_expressions( + /// &self, + /// f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + /// ) -> Result<TreeNodeRecursion> { + /// apply_expression_roots([&self.predicate], f) + /// } + /// ``` + /// + /// ## Node with no expressions (e.g., EmptyExec, MemoryExec) + /// ```ignore + /// fn apply_expressions( + /// &self, + /// _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + /// ) -> Result<TreeNodeRecursion> { + /// Ok(TreeNodeRecursion::Continue) + /// } + /// ``` + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion>; + + /// Deprecated. + /// + /// DataFusion will remove this method in the future in favor of + /// [`ExecutionPlan::replace_children`]. + /// + /// Note that this method is still required by the trait; implementations + /// should delegate to [`ExecutionPlan::replace_children`] with + /// [`ChildrenPropertiesMode::Recompute`]. + /// + /// # Example Implementation + /// ``` + /// # #![allow(deprecated)] + /// # use std::fmt; + /// # use std::sync::Arc; + /// # use datafusion_common::Result; + /// # use datafusion_common::tree_node::TreeNodeRecursion; + /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + /// # use datafusion_physical_expr::PhysicalExpr; + /// # use datafusion_physical_plan::{ + /// # ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + /// # PlanProperties, ReplaceChildrenOptions, + /// # }; + /// # #[derive(Debug)] + /// # struct MyExec { + /// # input: Arc<dyn ExecutionPlan>, + /// # } + /// # impl DisplayAs for MyExec { + /// # fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + /// # write!(f, "MyExec") + /// # } + /// # } + /// impl ExecutionPlan for MyExec { + /// // ... + /// # fn name(&self) -> &'static str { + /// # "MyExec" + /// # } + /// # fn properties(&self) -> &Arc<PlanProperties> { + /// # self.input.properties() + /// # } + /// # fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + /// # vec![&self.input] + /// # } + /// # fn apply_expressions( + /// # &self, + /// # _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + /// # ) -> Result<TreeNodeRecursion> { + /// # Ok(TreeNodeRecursion::Continue) + /// # } + /// # fn execute( + /// # &self, + /// # _partition: usize, + /// # _context: Arc<TaskContext>, + /// # ) -> Result<SendableRecordBatchStream> { + /// # unimplemented!() + /// # } + /// fn replace_children( + /// self: Arc<Self>, + /// mut children: Vec<Arc<dyn ExecutionPlan>>, + /// _options: ReplaceChildrenOptions, + /// ) -> Result<Arc<dyn ExecutionPlan>> { + /// Ok(Arc::new(MyExec { + /// input: children.swap_remove(0), + /// })) + /// } + /// + /// fn with_new_children( + /// self: Arc<Self>, + /// children: Vec<Arc<dyn ExecutionPlan>>, + /// ) -> Result<Arc<dyn ExecutionPlan>> { + /// // call into `replace_children` with `ReplaceChildrenOptions` + /// self.replace_children( + /// children, + /// ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + /// ) + /// } + /// } + /// ``` + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`" + )] + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>>; + + /// Deprecated. Implement [`ExecutionPlan::replace_children`] instead. + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`" + )] + #[expect(deprecated)] + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.with_new_children(children) + } + + /// Reset any internal state within this [`ExecutionPlan`]. + /// + /// This method is called when an [`ExecutionPlan`] needs to be re-executed, + /// such as in recursive queries. Unlike [`ExecutionPlan::replace_children`], this method + /// ensures that any stateful components (e.g., [`DynamicFilterPhysicalExpr`]) + /// are reset to their initial state. + /// + /// The default implementation simply calls [`ExecutionPlan::replace_children`] with the existing children, + /// effectively creating a new instance of the [`ExecutionPlan`] with the same children but without + /// necessarily resetting any internal state. Implementations that require resetting of some + /// internal state should override this method to provide the necessary logic. + /// + /// This method should *not* reset state recursively for children, as it is expected that + /// it will be called from within a walk of the execution plan tree so that it will be called on each child later + /// or was already called on each child. + /// + /// Note to implementers: unlike [`ExecutionPlan::replace_children`] this method does not accept new children as an argument, + /// thus it is expected that any cached plan properties will remain valid after the reset. + /// + /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + let children = self.children().into_iter().cloned().collect(); + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + /// If supported, attempt to increase the partitioning of this `ExecutionPlan` to + /// produce `target_partitions` partitions. + /// + /// If the `ExecutionPlan` does not support changing its partitioning, + /// returns `Ok(None)` (the default). + /// + /// If the `ExecutionPlan` can increase its partitioning, but not to + /// `target_partitions`, it may return an ExecutionPlan with fewer + /// partitions. This might happen, for example, if each new partition would + /// be too small to be efficiently processed individually. + /// + /// The DataFusion optimizer attempts to use as many threads as possible by + /// repartitioning its inputs to match the target number of threads + /// available (`target_partitions`). Some data sources, such as the built in + /// CSV and Parquet readers, implement this method as they are able to read + /// from their input files in parallel, regardless of how the source data is + /// split amongst files. + fn repartitioned( + &self, + _target_partitions: usize, + _config: &ConfigOptions, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + Ok(None) + } + + /// Begin execution of `partition`, returning a [`Stream`] of + /// [`RecordBatch`]es. + /// + /// # Notes + /// + /// The `execute` method itself is not `async` but it returns an `async` + /// [`futures::stream::Stream`]. This `Stream` should incrementally compute + /// the output, `RecordBatch` by `RecordBatch` (in a streaming fashion). + /// Most `ExecutionPlan`s should not do any work before the first + /// `RecordBatch` is requested from the stream. + /// + /// [`RecordBatchStreamAdapter`] can be used to convert an `async` + /// [`Stream`] into a [`SendableRecordBatchStream`]. + /// + /// Using `async` `Streams` allows for network I/O during execution and + /// takes advantage of Rust's built in support for `async` continuations and + /// crate ecosystem. + /// + /// [`Stream`]: futures::stream::Stream + /// [`StreamExt`]: futures::stream::StreamExt + /// [`TryStreamExt`]: futures::stream::TryStreamExt + /// [`RecordBatchStreamAdapter`]: crate::stream::RecordBatchStreamAdapter + /// + /// # Error handling + /// + /// Any error that occurs during execution is sent as an `Err` in the output + /// stream. + /// + /// `ExecutionPlan` implementations in DataFusion cancel additional work + /// immediately once an error occurs. The rationale is that if the overall + /// query will return an error, any additional work such as continued + /// polling of inputs will be wasted as it will be thrown away. + /// + /// # Cancellation / Aborting Execution + /// + /// The [`Stream`] that is returned must ensure that any allocated resources + /// are freed when the stream itself is dropped. This is particularly + /// important for [`spawn`]ed tasks or threads. Unless care is taken to + /// "abort" such tasks, they may continue to consume resources even after + /// the plan is dropped, generating intermediate results that are never + /// used. + /// Thus, [`spawn`] is disallowed, and instead use [`SpawnedTask`]. + /// + /// To enable timely cancellation, the [`Stream`] that is returned must not + /// block the CPU indefinitely and must yield back to the tokio runtime regularly. + /// In a typical [`ExecutionPlan`], this automatically happens unless there are + /// special circumstances; e.g. when the computational complexity of processing a + /// batch is superlinear. See this [general guideline][async-guideline] for more context + /// on this point, which explains why one should avoid spending a long time without + /// reaching an `await`/yield point in asynchronous runtimes. + /// This can be achieved by using the utilities from the [`coop`](crate::coop) module, by + /// manually returning [`Poll::Pending`] and setting up wakers appropriately, or by calling + /// [`tokio::task::yield_now()`] when appropriate. + /// In special cases that warrant manual yielding, determination for "regularly" may be + /// made using the [Tokio task budget](https://docs.rs/tokio/latest/tokio/task/coop/index.html), + /// a timer (being careful with the overhead-heavy system call needed to take the time), or by + /// counting rows or batches. + /// + /// The [cancellation benchmark] tracks some cases of how quickly queries can + /// be cancelled. + /// + /// For more details see [`SpawnedTask`], [`JoinSet`] and [`RecordBatchReceiverStreamBuilder`] + /// for structures to help ensure all background tasks are cancelled. + /// + /// [`spawn`]: tokio::task::spawn + /// [cancellation benchmark]: https://github.com/apache/datafusion/blob/main/benchmarks/README.md#cancellation + /// [`JoinSet`]: datafusion_common_runtime::JoinSet + /// [`SpawnedTask`]: datafusion_common_runtime::SpawnedTask + /// [`RecordBatchReceiverStreamBuilder`]: crate::stream::RecordBatchReceiverStreamBuilder + /// [`Poll::Pending`]: std::task::Poll::Pending + /// [async-guideline]: https://ryhl.io/blog/async-what-is-blocking/ + /// + /// # Implementation Examples + /// + /// While `async` `Stream`s have a non trivial learning curve, the + /// [`futures`] crate provides [`StreamExt`] and [`TryStreamExt`] + /// which help simplify many common operations. + /// + /// Here are some common patterns: + /// + /// ## Return Precomputed `RecordBatch` + /// + /// We can return a precomputed `RecordBatch` as a `Stream`: + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow::array::RecordBatch; + /// # use arrow::datatypes::SchemaRef; + /// # use datafusion_common::Result; + /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + /// # use datafusion_physical_plan::memory::MemoryStream; + /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + /// struct MyPlan { + /// batch: RecordBatch, + /// } + /// + /// impl MyPlan { + /// fn execute( + /// &self, + /// partition: usize, + /// context: Arc<TaskContext>, + /// ) -> Result<SendableRecordBatchStream> { + /// // use functions from futures crate to convert the batch into a stream + /// let fut = futures::future::ready(Ok(self.batch.clone())); + /// let stream = futures::stream::once(fut); + /// Ok(Box::pin(RecordBatchStreamAdapter::new( + /// self.batch.schema(), + /// stream, + /// ))) + /// } + /// } + /// ``` + /// + /// ## Lazily (async) Compute `RecordBatch` + /// + /// We can also lazily compute a `RecordBatch` when the returned `Stream` is polled + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow::array::RecordBatch; + /// # use arrow::datatypes::SchemaRef; + /// # use datafusion_common::Result; + /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + /// # use datafusion_physical_plan::memory::MemoryStream; + /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + /// struct MyPlan { + /// schema: SchemaRef, + /// } + /// + /// /// Returns a single batch when the returned stream is polled + /// async fn get_batch() -> Result<RecordBatch> { + /// todo!() + /// } + /// + /// impl MyPlan { + /// fn execute( + /// &self, + /// partition: usize, + /// context: Arc<TaskContext>, + /// ) -> Result<SendableRecordBatchStream> { + /// let fut = get_batch(); + /// let stream = futures::stream::once(fut); + /// Ok(Box::pin(RecordBatchStreamAdapter::new( + /// self.schema.clone(), + /// stream, + /// ))) + /// } + /// } + /// ``` + /// + /// ## Lazily (async) create a Stream + /// + /// If you need to create the return `Stream` using an `async` function, + /// you can do so by flattening the result: + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow::array::RecordBatch; + /// # use arrow::datatypes::SchemaRef; + /// # use futures::TryStreamExt; + /// # use datafusion_common::Result; + /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + /// # use datafusion_physical_plan::memory::MemoryStream; + /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + /// struct MyPlan { + /// schema: SchemaRef, + /// } + /// + /// /// async function that returns a stream + /// async fn get_batch_stream() -> Result<SendableRecordBatchStream> { + /// todo!() + /// } + /// + /// impl MyPlan { + /// fn execute( + /// &self, + /// partition: usize, + /// context: Arc<TaskContext>, + /// ) -> Result<SendableRecordBatchStream> { + /// // A future that yields a stream + /// let fut = get_batch_stream(); + /// // Use TryStreamExt::try_flatten to flatten the stream of streams + /// let stream = futures::stream::once(fut).try_flatten(); + /// Ok(Box::pin(RecordBatchStreamAdapter::new( + /// self.schema.clone(), + /// stream, + /// ))) + /// } + /// } + /// ``` + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream>; + + /// Return a snapshot of the set of [`Metric`]s for this + /// [`ExecutionPlan`]. If no `Metric`s are available, return None. + /// + /// While the values of the metrics in the returned + /// [`MetricsSet`]s may change as execution progresses, the + /// specific metrics will not. + /// + /// Once `self.execute()` has returned (technically the future is + /// resolved) for all available partitions, the set of metrics + /// should be complete. If this function is called prior to + /// `execute()` new metrics may appear in subsequent calls. + fn metrics(&self) -> Option<MetricsSet> { + None + } + + /// Returns statistics for a specific partition of this `ExecutionPlan` node. + /// + /// Deprecated: use [`StatisticsContext::compute`] instead. + /// + /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute + #[deprecated(since = "55.0.0", note = "Use StatisticsContext::compute instead")] + fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> { + if let Some(idx) = partition { + // Validate partition index + let partition_count = self.properties().partitioning.partition_count(); + assert_or_internal_err!( + idx < partition_count, + "Invalid partition index: {}, the partition count is {}", + idx, + partition_count + ); + } + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } + + /// Returns statistics for a specific partition of this `ExecutionPlan` node, + /// given pre-computed child statistics. + /// + /// If statistics are not available, should return [`Statistics::new_unknown`] + /// (the default), not an error. + /// If `args.partition()` is `None`, it returns statistics for all partitions. + /// + /// Implementations should not call [`StatisticsContext::compute`] from within + /// this method; child statistics are provided via `input_stats`. + /// + /// Use [`StatisticsContext::compute`] to initiate a full plan-tree walk. + /// + /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + #[expect(deprecated)] + self.partition_statistics(args.partition()) + } + + /// Returns, per child, which statistics the [`StatisticsContext`] should resolve + /// before calling [`Self::statistics_from_inputs`]. + /// + /// One entry per child (same order as [`Self::children`]): [`ChildStats::At`] + /// requests the child's statistics at a partition (`None` = overall); + /// [`ChildStats::Skip`] omits a child whose statistics this node does not need + /// (a `Statistics::new_unknown` placeholder fills its `input_stats` slot). + /// + /// The default skips every child, so a node that derives nothing from its + /// children (for example one that only overrides the deprecated + /// [`Self::partition_statistics`]) triggers no child traversal. A node that reads + /// `input_stats` in [`Self::statistics_from_inputs`] must override this to declare + /// the children it uses. + /// + /// [`StatisticsContext`]: crate::statistics::StatisticsContext + fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> { + self.children().iter().map(|_| ChildStats::Skip).collect() + } + + /// Returns `true` if a limit can be safely pushed down through this + /// `ExecutionPlan` node. + /// + /// If this method returns `true`, and the query plan contains a limit at + /// the output of this node, DataFusion will push the limit to the input + /// of this node. + fn supports_limit_pushdown(&self) -> bool { + false + } + + /// Returns a fetching variant of this `ExecutionPlan` node, if it supports + /// fetch limits. Returns `None` otherwise. + /// + /// See physical optimizer rule [`limit_pushdown`] for details. + /// + /// [`limit_pushdown`]: https://docs.rs/datafusion/latest/datafusion/physical_optimizer/limit_pushdown/index.html + fn with_fetch(&self, _limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + None + } + + /// Gets the fetch count for the operator, `None` means there is no fetch. + fn fetch(&self) -> Option<usize> { + None + } + + /// Gets the effect on cardinality, if known + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Unknown + } + + /// Attempts to push down the given projection into the input of this `ExecutionPlan`. + /// + /// If the operator supports this optimization, the resulting plan will be: + /// `self_new <- projection <- source`, starting from `projection <- self <- source`. + /// Otherwise, it returns the current `ExecutionPlan` as-is. + /// + /// Returns `Ok(Some(...))` if pushdown is applied, `Ok(None)` if it is not supported + /// or not possible, or `Err` on failure. + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + Ok(None) + } + + /// Collect filters that this node can push down to its children. + /// Filters that are being pushed down from parents are passed in, + /// and the node may generate additional filters to push down. + /// For example, given the plan FilterExec -> HashJoinExec -> DataSourceExec, + /// what will happen is that we recurse down the plan calling `ExecutionPlan::gather_filters_for_pushdown`: + /// 1. `FilterExec::gather_filters_for_pushdown` is called with no parent + /// filters so it only returns that `FilterExec` wants to push down its own predicate. + /// 2. `HashJoinExec::gather_filters_for_pushdown` is called with the filter from + /// `FilterExec`, which it only allows to push down to one side of the join (unless it's on the join key) + /// but it also adds its own filters (e.g. pushing down a bloom filter of the hash table to the scan side of the join). + /// 3. `DataSourceExec::gather_filters_for_pushdown` is called with both filters from `HashJoinExec` + /// and `FilterExec`, however `DataSourceExec::gather_filters_for_pushdown` doesn't actually do anything + /// since it has no children and no additional filters to push down. + /// It's only once [`ExecutionPlan::handle_child_pushdown_result`] is called on `DataSourceExec` as we recurse + /// up the plan that `DataSourceExec` can actually bind the filters. + /// + /// The default implementation bars all parent filters from being pushed down and adds no new filters. + /// This is the safest option, making filter pushdown opt-in on a per-node basis. + /// + /// There are two different phases in filter pushdown, which some operators may handle the same and some differently. + /// Depending on the phase the operator may or may not be allowed to modify the plan. + /// See [`FilterPushdownPhase`] for more details. + /// + /// Implementations must preserve the order of `parent_filters` in the + /// returned child [`FilterDescription`]: each child parent-filter result is + /// matched back to the corresponding input parent filter by position. + /// Unsupported filters should therefore be marked unsupported in place, + /// rather than removed or appended after supported filters. + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + Ok(FilterDescription::all_unsupported( + &parent_filters, + &self.children(), + )) + } + + /// Handle the result of a child pushdown. + /// + /// This method is called as we recurse back up the plan tree after pushing + /// filters down to child nodes via [`ExecutionPlan::gather_filters_for_pushdown`]. + /// It allows the current node to process the results of filter pushdown from + /// its children, deciding whether to absorb filters, modify the plan, or pass + /// filters back up to its parent. + /// + /// **Purpose and Context:** + /// Filter pushdown is a critical optimization in DataFusion that aims to + /// reduce the amount of data processed by applying filters as early as + /// possible in the query plan. This method is part of the second phase of + /// filter pushdown, where results are propagated back up the tree after + /// being pushed down. Each node can inspect the pushdown results from its + /// children and decide how to handle any unapplied filters, potentially + /// optimizing the plan structure or filter application. + /// + /// **Behavior in Different Nodes:** + /// - For a `DataSourceExec`, this often means absorbing the filters to apply + /// them during the scan phase (late materialization), reducing the data + /// read from the source. + /// - A `FilterExec` may absorb any filters its children could not handle, + /// combining them with its own predicate. If no filters remain (i.e., the + /// predicate becomes trivially true), it may remove itself from the plan + /// altogether. It typically marks parent filters as supported, indicating + /// they have been handled. + /// - A `HashJoinExec` might ignore the pushdown result if filters need to + /// be applied during the join operation. It passes the parent filters back + /// up wrapped in [`FilterPushdownPropagation::if_any`], discarding + /// any self-filters from children. + /// + /// **Example Walkthrough:** + /// Consider a query plan: `FilterExec (f1) -> HashJoinExec -> DataSourceExec`. + /// 1. **Downward Phase (`gather_filters_for_pushdown`):** Starting at + /// `FilterExec`, the filter `f1` is gathered and pushed down to + /// `HashJoinExec`. `HashJoinExec` may allow `f1` to pass to one side of + /// the join or add its own filters (e.g., a min-max filter from the build side), + /// then pushes filters to `DataSourceExec`. `DataSourceExec`, being a leaf node, + /// has no children to push to, so it prepares to handle filters in the + /// upward phase. + /// 2. **Upward Phase (`handle_child_pushdown_result`):** Starting at + /// `DataSourceExec`, it absorbs applicable filters from `HashJoinExec` + /// for late materialization during scanning, marking them as supported. + /// `HashJoinExec` receives the result, decides whether to apply any + /// remaining filters during the join, and passes unhandled filters back + /// up to `FilterExec`. `FilterExec` absorbs any unhandled filters, + /// updates its predicate if necessary, or removes itself if the predicate + /// becomes trivial (e.g., `lit(true)`), and marks filters as supported + /// for its parent. + /// + /// The default implementation is a no-op that passes the result of pushdown + /// from the children to its parent transparently, ensuring no filters are + /// lost if a node does not override this behavior. + /// + /// **Notes for Implementation:** + /// When returning filters via [`FilterPushdownPropagation`], the order of + /// filters need not match the order they were passed in via + /// `child_pushdown_result`. However, preserving the order is recommended for + /// debugging and ease of reasoning about the resulting plans. + /// + /// **Helper Methods for Customization:** + /// There are various helper methods to simplify implementing this method: + /// - [`FilterPushdownPropagation::if_any`]: Marks all parent filters as + /// supported as long as at least one child supports them. + /// - [`FilterPushdownPropagation::if_all`]: Marks all parent filters as + /// supported as long as all children support them. + /// - [`FilterPushdownPropagation::with_parent_pushdown_result`]: Allows adding filters + /// to the propagation result, indicating which filters are supported by + /// the current node. + /// - [`FilterPushdownPropagation::with_updated_node`]: Allows updating the + /// current node in the propagation result, used if the node + /// has modified its plan based on the pushdown results. + /// + /// **Filter Pushdown Phases:** + /// There are two different phases in filter pushdown (`Pre` and others), + /// which some operators may handle differently. Depending on the phase, the + /// operator may or may not be allowed to modify the plan. See + /// [`FilterPushdownPhase`] for more details on phase-specific behavior. + /// + /// [`PushedDownPredicate::supported`]: crate::filter_pushdown::PushedDownPredicate::supported + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + Ok(FilterPushdownPropagation::if_all(child_pushdown_result)) + } + + /// Injects arbitrary run-time state into this execution plan, returning a new plan + /// instance that incorporates that state *if* it is relevant to the concrete + /// node implementation. + /// + /// This is a generic entry point: the `state` can be any type wrapped in + /// `Arc<dyn Any + Send + Sync>`. A node that cares about the state should + /// down-cast it to the concrete type it expects and, if successful, return a + /// modified copy of itself that captures the provided value. If the state is + /// not applicable, the default behaviour is to return `None` so that parent + /// nodes can continue propagating the attempt further down the plan tree. + /// + /// For example, [`WorkTableExec`](crate::work_table::WorkTableExec) + /// down-casts the supplied state to an `Arc<WorkTable>` + /// in order to wire up the working table used during recursive-CTE execution. + /// Similar patterns can be followed by custom nodes that need late-bound + /// dependencies or shared state. + fn with_new_state( + &self, + _state: Arc<dyn Any + Send + Sync>, + ) -> Option<Arc<dyn ExecutionPlan>> { + None + } + + /// Try to push down sort ordering requirements to this node. + /// + /// This method is called during sort pushdown optimization to determine if this + /// node can optimize for a requested sort ordering. Implementations should: + /// + /// - Return [`SortOrderPushdownResult::Exact`] if the node can guarantee the exact + /// ordering (allowing the Sort operator to be removed) + /// - Return [`SortOrderPushdownResult::Inexact`] if the node can optimize for the + /// ordering but cannot guarantee perfect sorting (Sort operator is kept) + /// - Return [`SortOrderPushdownResult::Unsupported`] if the node cannot optimize + /// for the ordering + /// + /// For transparent nodes (that preserve ordering), implement this to delegate to + /// children and wrap the result with a new instance of this node. + /// + /// Default implementation returns `Unsupported`. + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + Ok(SortOrderPushdownResult::Unsupported) + } + + /// Returns a variant of this `ExecutionPlan` that is aware of order-sensitivity. + /// + /// This is used to signal to data sources that the output ordering must be + /// preserved, even if it might be more efficient to ignore it (e.g. by + /// skipping some row groups in Parquet). + /// + fn with_preserve_order( + &self, + _preserve_order: bool, + ) -> Option<Arc<dyn ExecutionPlan>> { + None + } + + /// Serialize this plan to its protobuf representation, if it knows how. + /// + /// This is the `ExecutionPlan` analog of + /// [`PhysicalExpr::try_to_proto`]. + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller + /// (`datafusion-proto`) falls back to the central downcast chain. Every + /// un-migrated plan keeps its existing behavior. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// * `Err(_)` — a real failure (e.g. a child failed to serialize). + /// + /// Only *self-contained* plans should override this — see [`crate::proto`] + /// for the session-dependency boundary. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + Ok(None) + } +} + +/// Options for [`ExecutionPlan::replace_children`] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ReplaceChildrenOptions { + /// Describes how plan properties should be handled for the replacement + /// children. + pub children_properties: ChildrenPropertiesMode, +} + +impl ReplaceChildrenOptions { + /// Create new options for [`ExecutionPlan::replace_children`]. + pub const fn new(children_properties: ChildrenPropertiesMode) -> Self { + Self { + children_properties, + } + } +} + +/// Indicates whether the plan properties of the new children must be recomputed. +/// +/// Part of [`ReplaceChildrenOptions`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildrenPropertiesMode { + /// The plan properties of the new children are identical to the properties + /// of the existing children, so we can skip recomputation. + Keep, + /// The plan properties of the new children are different from the properties + /// of the existing children, so we must recompute the properties from scratch. + Recompute, +} + +/// Allows a type to be treated as a reference to an +/// [`Arc<dyn PhysicalExpr>`]. +/// +/// Used by [`apply_expression_roots`]. +pub trait AsPhysicalExprRef { + /// Returns the referenced physical expression. + fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr>; +} + +/// Allows an [`Arc<dyn PhysicalExpr>`] to be treated as a reference to itself. +/// +/// This is needed because `Arc<dyn PhysicalExpr>` does not implement +/// `AsRef<Arc<dyn PhysicalExpr>>`. +impl AsPhysicalExprRef for Arc<dyn PhysicalExpr> { + fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> { + self + } +} + +/// Allows a [`ProjectionExpr`] to be treated as a reference to its +/// [`Arc<dyn PhysicalExpr>`]. +impl AsPhysicalExprRef for ProjectionExpr { + fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> { + self.as_ref() + } +} + +impl<T> AsPhysicalExprRef for &T +where + T: AsPhysicalExprRef + ?Sized, +{ + fn as_physical_expr_ref(&self) -> &Arc<dyn PhysicalExpr> { + (*self).as_physical_expr_ref() + } +} + +/// Applies `f` to a shallow sequence of physical expression roots. +/// +/// [`TreeNodeRecursion::Stop`] stops iteration and is returned immediately. +/// [`TreeNodeRecursion::Jump`] is normalized to [`TreeNodeRecursion::Continue`] +/// because this function does not visit expression children. +pub fn apply_expression_roots<I>( + roots: I, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, +) -> Result<TreeNodeRecursion> +where + I: IntoIterator, + I::Item: AsPhysicalExprRef, +{ + for root in roots { + match f(root.as_physical_expr_ref())? { + TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop), + TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {} + } + } + Ok(TreeNodeRecursion::Continue) +} + +/// Returns whether `plan` contains a physical expression with `expression_id`. +/// +/// This traverses both the execution plan and the children of each expression root +/// reported by [`ExecutionPlan::apply_expressions`]. +pub(crate) fn plan_contains_expression_id( + plan: &Arc<dyn ExecutionPlan>, + expression_id: u64, +) -> Result<bool> { + let mut found = false; + plan.apply(|node| { + node.apply_expressions(&mut |root| { + root.apply(|expr| { + if expr.expression_id() == Some(expression_id) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + Ok(found) +} + +impl dyn ExecutionPlan { + /// Returns `true` if the plan is of type `T`. + /// + /// If this plan provides a [`ExecutionPlan::downcast_delegate`], delegates + /// to it. + /// + /// Prefer this over `downcast_ref::<T>().is_some()`. Works correctly when + /// called on `Arc<dyn ExecutionPlan>` via auto-deref. + pub fn is<T: ExecutionPlan>(&self) -> bool { + match self.downcast_delegate() { + Some(delegate) => delegate.is::<T>(), + None => (self as &dyn Any).is::<T>(), + } + } + + /// Attempts to downcast this plan to a concrete type `T`, returning `None` + /// if the plan is not of that type. + /// + /// If this plan provides a [`ExecutionPlan::downcast_delegate`], delegates + /// to it. + /// + /// Works correctly when called on `Arc<dyn ExecutionPlan>` via auto-deref, + /// unlike `(&arc as &dyn Any).downcast_ref::<T>()` which would attempt to + /// downcast the `Arc` itself. + pub fn downcast_ref<T: ExecutionPlan>(&self) -> Option<&T> { + match self.downcast_delegate() { + Some(delegate) => delegate.downcast_ref::<T>(), + None => (self as &dyn Any).downcast_ref(), + } + } +} + +/// [`ExecutionPlan`] Invariant Level +/// +/// What set of assertions ([Invariant]s) holds for a particular `ExecutionPlan` +/// +/// [Invariant]: https://en.wikipedia.org/wiki/Invariant_(mathematics)#Invariants_in_computer_science +#[derive(Clone, Copy)] +pub enum InvariantLevel { + /// Invariants that are always true for the [`ExecutionPlan`] node + /// such as the number of expected children. + Always, + /// Invariants that must hold true for the [`ExecutionPlan`] node + /// to be "executable", such as ordering and/or distribution requirements + /// being fulfilled. + Executable, +} + +/// Extension trait provides an easy API to fetch various properties of +/// [`ExecutionPlan`] objects based on [`ExecutionPlan::properties`]. +pub trait ExecutionPlanProperties { + /// Specifies how the output of this `ExecutionPlan` is split into + /// partitions. + fn output_partitioning(&self) -> &Partitioning; + + /// If the output of this `ExecutionPlan` within each partition is sorted, + /// returns `Some(keys)` describing the ordering. A `None` return value + /// indicates no assumptions should be made on the output ordering. + /// + /// For example, `SortExec` (obviously) produces sorted output as does + /// `SortPreservingMergeStream`. Less obviously, `Projection` produces sorted + /// output if its input is sorted as it does not reorder the input rows. + fn output_ordering(&self) -> Option<&LexOrdering>; + + /// Boundedness information of the stream corresponding to this `ExecutionPlan`. + /// For more details, see [`Boundedness`]. + fn boundedness(&self) -> Boundedness; + + /// Indicates how the stream of this `ExecutionPlan` emits its results. + /// For more details, see [`EmissionType`]. + fn pipeline_behavior(&self) -> EmissionType; + + /// Get the [`EquivalenceProperties`] within the plan. + /// + /// Equivalence properties tell DataFusion what columns are known to be + /// equal, during various optimization passes. By default, this returns "no + /// known equivalences" which is always correct, but may cause DataFusion to + /// unnecessarily resort data. + /// + /// If this ExecutionPlan makes no changes to the schema of the rows flowing + /// through it or how columns within each row relate to each other, it + /// should return the equivalence properties of its input. For + /// example, since [`FilterExec`] may remove rows from its input, but does not + /// otherwise modify them, it preserves its input equivalence properties. + /// However, since `ProjectionExec` may calculate derived expressions, it + /// needs special handling. + /// + /// See also [`ExecutionPlan::maintains_input_order`] and [`Self::output_ordering`] + /// for related concepts. + /// + /// [`FilterExec`]: crate::filter::FilterExec + fn equivalence_properties(&self) -> &EquivalenceProperties; +} + +impl ExecutionPlanProperties for Arc<dyn ExecutionPlan> { + fn output_partitioning(&self) -> &Partitioning { + self.properties().output_partitioning() + } + + fn output_ordering(&self) -> Option<&LexOrdering> { + self.properties().output_ordering() + } + + fn boundedness(&self) -> Boundedness { + self.properties().boundedness + } + + fn pipeline_behavior(&self) -> EmissionType { + self.properties().emission_type + } + + fn equivalence_properties(&self) -> &EquivalenceProperties { + self.properties().equivalence_properties() + } +} + +impl ExecutionPlanProperties for &dyn ExecutionPlan { + fn output_partitioning(&self) -> &Partitioning { + self.properties().output_partitioning() + } + + fn output_ordering(&self) -> Option<&LexOrdering> { + self.properties().output_ordering() + } + + fn boundedness(&self) -> Boundedness { + self.properties().boundedness + } + + fn pipeline_behavior(&self) -> EmissionType { + self.properties().emission_type + } + + fn equivalence_properties(&self) -> &EquivalenceProperties { + self.properties().equivalence_properties() + } +} + +/// Represents whether a stream of data **generated** by an operator is bounded (finite) +/// or unbounded (infinite). +/// +/// This is used to determine whether an execution plan will eventually complete +/// processing all its data (bounded) or could potentially run forever (unbounded). +/// +/// For unbounded streams, it also tracks whether the operator requires finite memory +/// to process the stream or if memory usage could grow unbounded. +/// +/// Boundedness of the output stream is based on the boundedness of the input stream and the nature of +/// the operator. For example, limit or topk with fetch operator can convert an unbounded stream to a bounded stream. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Boundedness { + /// The data stream is bounded (finite) and will eventually complete + Bounded, + /// The data stream is unbounded (infinite) and could run forever + Unbounded { + /// Whether this operator requires infinite memory to process the unbounded stream. + /// If false, the operator can process an infinite stream with bounded memory. + /// If true, memory usage may grow unbounded while processing the stream. + /// + /// For example, `Median` requires infinite memory to compute the median of an unbounded stream. + /// `Min/Max` requires infinite memory if the stream is unordered, but can be computed with bounded memory if the stream is ordered. + requires_infinite_memory: bool, + }, +} + +impl Boundedness { + pub fn is_unbounded(&self) -> bool { + matches!(self, Boundedness::Unbounded { .. }) + } +} + +/// Represents how an operator emits its output records. +/// +/// This is used to determine whether an operator emits records incrementally as they arrive, +/// only emits a final result at the end, or can do both. Note that it generates the output -- record batch with `batch_size` rows +/// but it may still buffer data internally until it has enough data to emit a record batch or the source is exhausted. +/// +/// For example, in the following plan: +/// ```text +/// SortExec [EmissionType::Final] +/// |_ on: [col1 ASC] +/// FilterExec [EmissionType::Incremental] +/// |_ pred: col2 > 100 +/// DataSourceExec [EmissionType::Incremental] +/// |_ file: "data.csv" +/// ``` +/// - DataSourceExec emits records incrementally as it reads from the file +/// - FilterExec processes and emits filtered records incrementally as they arrive +/// - SortExec must wait for all input records before it can emit the sorted result, +/// since it needs to see all values to determine their final order +/// +/// Left joins can emit both incrementally and finally: +/// - Incrementally emit matches as they are found +/// - Finally emit non-matches after all input is processed +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EmissionType { + /// Records are emitted incrementally as they arrive and are processed + Incremental, + /// Records are only emitted once all input has been processed + Final, + /// Records can be emitted both incrementally and as a final result + Both, +} + +/// Represents whether an operator's `Stream` has been implemented to actively cooperate with the +/// Tokio scheduler or not. Please refer to the [`coop`](crate::coop) module for more details. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SchedulingType { + /// The stream generated by [`execute`](ExecutionPlan::execute) does not actively participate in + /// cooperative scheduling. This means the implementation of the `Stream` returned by + /// [`ExecutionPlan::execute`] does not contain explicit task budget consumption such as + /// [`tokio::task::coop::consume_budget`]. + /// + /// `NonCooperative` is the default value and is acceptable for most operators. Please refer to + /// the [`coop`](crate::coop) module for details on when it may be useful to use + /// `Cooperative` instead. + NonCooperative, + /// The stream generated by [`execute`](ExecutionPlan::execute) actively participates in + /// cooperative scheduling by consuming task budget when it was able to produce a + /// [`RecordBatch`]. + Cooperative, +} + +/// Represents how an operator's stream drives [`RecordBatch`] production +/// relative to downstream demand. +/// +/// This is execution-topology metadata for optimizers. It distinguishes streams +/// whose batch production is driven directly by downstream calls to +/// `Stream::poll_next` from streams that may also drive input or output +/// production independently, such as by spawning tasks or buffering batches +/// ahead of demand. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EvaluationType { + /// The stream generated by [`execute`](ExecutionPlan::execute) is + /// demand-driven: it produces [`RecordBatch`]es in response to downstream + /// calls to `Stream::poll_next`. + /// + /// Filter, projection, and join operators are examples of lazy operators. + /// + /// Lazy operators are also known as demand-driven operators. + Lazy, + /// The stream generated by [`execute`](ExecutionPlan::execute) may drive + /// input or output [`RecordBatch`] production ahead of, or independently + /// from, downstream calls to `Stream::poll_next`. + /// + /// Eager operators commonly poll input streams from spawned Tokio tasks, + /// buffer batches ahead of demand, or otherwise create an independent + /// child-polling pipeline. Eager work may start when `execute` creates the + /// stream or when the returned stream is first polled; that timing is an + /// implementation detail. + /// + /// Repartition, coalesce partitions, sort-preserving merge, buffer, and + /// analyze operators are examples of eager operators. + /// + /// Eager operators are also known as a data-driven operators. + Eager, +} + +/// Utility to determine an operator's boundedness based on its children's boundedness. +/// +/// Assumes boundedness can be inferred from child operators: +/// - Unbounded (requires_infinite_memory: true) takes precedence. +/// - Unbounded (requires_infinite_memory: false) is considered next. +/// - Otherwise, the operator is bounded. +/// +/// **Note:** This is a general-purpose utility and may not apply to +/// all multi-child operators. Ensure your operator's behavior aligns +/// with these assumptions before using. +pub(crate) fn boundedness_from_children<'a>( + children: impl IntoIterator<Item = &'a Arc<dyn ExecutionPlan>>, +) -> Boundedness { + let mut unbounded_with_finite_mem = false; + + for child in children { + match child.boundedness() { + Boundedness::Unbounded { + requires_infinite_memory: true, + } => { + return Boundedness::Unbounded { + requires_infinite_memory: true, + }; + } + Boundedness::Unbounded { + requires_infinite_memory: false, + } => { + unbounded_with_finite_mem = true; + } + Boundedness::Bounded => {} + } + } + + if unbounded_with_finite_mem { + Boundedness::Unbounded { + requires_infinite_memory: false, + } + } else { + Boundedness::Bounded + } +} + +/// Determines the emission type of an operator based on its children's pipeline behavior. +/// +/// The precedence of emission types is: +/// - `Final` has the highest precedence. +/// - `Both` is next: if any child emits both incremental and final results, the parent inherits this behavior unless a `Final` is present. +/// - `Incremental` is the default if all children emit incremental results. +/// +/// **Note:** This is a general-purpose utility and may not apply to +/// all multi-child operators. Verify your operator's behavior aligns +/// with these assumptions. +pub(crate) fn emission_type_from_children<'a>( + children: impl IntoIterator<Item = &'a Arc<dyn ExecutionPlan>>, +) -> EmissionType { + let mut inc_and_final = false; + + for child in children { + match child.pipeline_behavior() { + EmissionType::Final => return EmissionType::Final, + EmissionType::Both => inc_and_final = true, + EmissionType::Incremental => continue, + } + } + + if inc_and_final { + EmissionType::Both + } else { + EmissionType::Incremental + } +} + +/// Stores plan properties used in query optimization. +/// +/// Serves as a cache for these properties, which are often +/// expensive to compute. +#[derive(Debug, Clone)] +pub struct PlanProperties { + /// See [ExecutionPlanProperties::equivalence_properties] + pub eq_properties: EquivalenceProperties, + /// See [ExecutionPlanProperties::output_partitioning] + pub partitioning: Partitioning, + /// See [ExecutionPlanProperties::pipeline_behavior] + pub emission_type: EmissionType, + /// See [ExecutionPlanProperties::boundedness] + pub boundedness: Boundedness, + pub evaluation_type: EvaluationType, + pub scheduling_type: SchedulingType, + /// See [ExecutionPlanProperties::output_ordering] + output_ordering: Option<LexOrdering>, +} + +impl PlanProperties { + /// Construct a new `PlanPropertiesCache` from the + pub fn new( + eq_properties: EquivalenceProperties, + partitioning: Partitioning, + emission_type: EmissionType, + boundedness: Boundedness, + ) -> Self { + // Output ordering can be derived from `eq_properties`. + let output_ordering = eq_properties.output_ordering(); + Self { + eq_properties, + partitioning, + emission_type, + boundedness, + evaluation_type: EvaluationType::Lazy, + scheduling_type: SchedulingType::NonCooperative, + output_ordering, + } + } + + /// Overwrite output partitioning with its new value. + pub fn with_partitioning(mut self, partitioning: Partitioning) -> Self { + self.partitioning = partitioning; + self + } + + /// Set equivalence properties having mut reference. + pub fn set_eq_properties(&mut self, eq_properties: EquivalenceProperties) { + // Changing equivalence properties also changes output ordering, so + // make sure to overwrite it: + self.output_ordering = eq_properties.output_ordering(); + self.eq_properties = eq_properties; + } + + /// Overwrite equivalence properties with its new value. + pub fn with_eq_properties(mut self, eq_properties: EquivalenceProperties) -> Self { + self.set_eq_properties(eq_properties); + self + } + + /// Overwrite boundedness with its new value. + pub fn with_boundedness(mut self, boundedness: Boundedness) -> Self { + self.boundedness = boundedness; + self + } + + /// Overwrite emission type with its new value. + pub fn with_emission_type(mut self, emission_type: EmissionType) -> Self { + self.emission_type = emission_type; + self + } + + /// Set the [`SchedulingType`]. + /// + /// Defaults to [`SchedulingType::NonCooperative`] + pub fn with_scheduling_type(mut self, scheduling_type: SchedulingType) -> Self { + self.scheduling_type = scheduling_type; + self + } + + /// Set the [`EvaluationType`]. + /// + /// Defaults to [`EvaluationType::Lazy`] + pub fn with_evaluation_type(mut self, drive_type: EvaluationType) -> Self { + self.evaluation_type = drive_type; + self + } + + /// Set constraints having mut reference. + pub fn set_constraints(&mut self, constraints: Constraints) { + self.eq_properties.set_constraints(constraints); + } + + /// Overwrite constraints with its new value. + pub fn with_constraints(mut self, constraints: Constraints) -> Self { + self.set_constraints(constraints); + self + } + + pub fn equivalence_properties(&self) -> &EquivalenceProperties { + &self.eq_properties + } + + pub fn output_partitioning(&self) -> &Partitioning { + &self.partitioning + } + + pub fn output_ordering(&self) -> Option<&LexOrdering> { + self.output_ordering.as_ref() + } + + /// Get schema of the node. + pub(crate) fn schema(&self) -> &SchemaRef { + self.eq_properties.schema() + } +} + +macro_rules! check_len { + ($target:expr, $func_name:ident, $expected_len:expr) => { + let actual_len = $target.$func_name().len(); + assert_eq_or_internal_err!( + actual_len, + $expected_len, + "{}::{} returned Vec with incorrect size: {} != {}", + $target.name(), + stringify!($func_name), + actual_len, + $expected_len + ); + }; +} + +/// All dynamic expressions must have an expression id. +fn check_dynamic_expression_invariants<P: ExecutionPlan + ?Sized>( + plan: &P, +) -> Result<()> { + let mut produced_ids = HashSet::new(); + for expr in plan.dynamic_expressions_produced() { + let Some(expression_id) = expr.expression_id() else { + return internal_err!( + "{}::dynamic_expressions_produced returned an expression without an expression ID", + plan.name() + ); + }; + assert_or_internal_err!( + produced_ids.insert(expression_id), + "{}::dynamic_expressions_produced returned duplicate expression ID {expression_id}", + plan.name() + ); + } + Ok(()) +} + +/// Checks a set of invariants that apply to all ExecutionPlan implementations. +/// Returns an error if the given node does not conform. +pub fn check_default_invariants<P: ExecutionPlan + ?Sized>( + plan: &P, + check: InvariantLevel, +) -> Result<(), DataFusionError> { + let children_len = plan.children().len(); + + check_len!(plan, maintains_input_order, children_len); + check_len!(plan, required_input_ordering, children_len); + check_len!(plan, benefits_from_input_partitioning, children_len); + plan.input_distribution_requirements() + .check_invariants(plan, check)?; + check_dynamic_expression_invariants(plan)?; + + Ok(()) +} + +/// Indicate whether a data exchange is needed for the input of `plan`. +/// +/// This identifies physical operators that redistribute child partitions or +/// gather multiple child partitions into one output partition: +/// +/// 1. RepartitionExec for non-round-robin repartitioning +/// 2. CoalescePartitionsExec for collapsing multiple partitions into one without ordering guarantee +/// 3. SortPreservingMergeExec for collapsing multiple sorted partitions into one with ordering guarantee +#[expect(clippy::needless_pass_by_value)] +pub fn need_data_exchange(plan: Arc<dyn ExecutionPlan>) -> bool { + if let Some(repartition) = plan.downcast_ref::<RepartitionExec>() { + !matches!(repartition.partitioning(), Partitioning::RoundRobinBatch(_)) + } else if let Some(coalesce) = plan.downcast_ref::<CoalescePartitionsExec>() { + coalesce.input().output_partitioning().partition_count() > 1 + } else if let Some(sort_preserving_merge) = + plan.downcast_ref::<SortPreservingMergeExec>() + { + sort_preserving_merge + .input() + .output_partitioning() + .partition_count() + > 1 + } else { + false + } +} + +/// Returns a plan with the given children, skipping as much work as possible. +/// +/// This helper is the single entry point for "rebuild a plan from new +/// children" and applies three layers of short-circuits, from cheapest to +/// most expensive: +/// +/// 1. **Same child pointers** — if every `children[i]` is `Arc::ptr_eq` to the +/// corresponding existing child, the original `plan` is returned +/// unchanged (no allocation, no [`ExecutionPlan::replace_children`] +/// call). +/// 2. **Same child properties** — if the children's `PlanProperties` Arcs +/// match (via [`has_same_children_properties`]), the plan's own +/// `PlanProperties` cache can be reused. This calls +/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Keep`], +/// which swaps the child pointers without recomputing `PlanProperties`. +/// 3. **Full recompute** — otherwise, delegate to +/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Recompute`], +/// which recomputes `PlanProperties` from scratch. +/// +/// The size of `children` must be equal to the size of `ExecutionPlan::children()`. +pub fn replace_children_if_necessary( + plan: Arc<dyn ExecutionPlan>, + children: Vec<Arc<dyn ExecutionPlan>>, +) -> Result<Arc<dyn ExecutionPlan>> { + let old_children = plan.children(); + assert_eq_or_internal_err!( + children.len(), + old_children.len(), + "Wrong number of children" + ); + if !children.is_empty() { + // Layer 1: same child pointers → return the plan unchanged. + if children + .iter() + .zip(old_children.iter()) + .all(|(c1, c2)| Arc::ptr_eq(c1, c2)) + { + return Ok(plan); + } + // Layer 2: same child properties → reuse `PlanProperties` cache. + if has_same_children_properties(plan.as_ref(), &children)? { + return plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ); + } + } + // Layer 3: full recompute. + plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) +} + +#[deprecated(since = "55.0.0", note = "Use `replace_children_if_necessary`")] +pub fn with_new_children_if_necessary( + plan: Arc<dyn ExecutionPlan>, + children: Vec<Arc<dyn ExecutionPlan>>, +) -> Result<Arc<dyn ExecutionPlan>> { + replace_children_if_necessary(plan, children) +} + +/// Return a [`DisplayableExecutionPlan`] wrapper around an +/// [`ExecutionPlan`] which can be displayed in various easier to +/// understand ways. +/// +/// See examples on [`DisplayableExecutionPlan`] +pub fn displayable(plan: &dyn ExecutionPlan) -> DisplayableExecutionPlan<'_> { + DisplayableExecutionPlan::new(plan) +} + +/// Execute the [ExecutionPlan] and collect the results in memory +pub async fn collect( + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, +) -> Result<Vec<RecordBatch>> { + let stream = execute_stream(plan, context)?; + crate::common::collect(stream).await +} + +/// Execute the [ExecutionPlan] and return a single stream of `RecordBatch`es. +/// +/// See [collect] to buffer the `RecordBatch`es in memory. +/// +/// # Aborting Execution +/// +/// Dropping the stream will abort the execution of the query, and free up +/// any allocated resources +#[expect( + clippy::needless_pass_by_value, + reason = "Public API that historically takes owned Arcs" +)] +pub fn execute_stream( + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, +) -> Result<SendableRecordBatchStream> { + match plan.output_partitioning().partition_count() { + 0 => Ok(Box::pin(EmptyRecordBatchStream::new(plan.schema()))), + 1 => plan.execute(0, context), + 2.. => { + // merge into a single partition + let plan = CoalescePartitionsExec::new(Arc::clone(&plan)); + // CoalescePartitionsExec must produce a single partition + assert_eq!(1, plan.properties().output_partitioning().partition_count()); + plan.execute(0, context) + } + } +} + +/// Execute the [ExecutionPlan] and collect the results in memory +pub async fn collect_partitioned( + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, +) -> Result<Vec<Vec<RecordBatch>>> { + // Avoid `JoinSet::spawn` for single partition + if plan.output_partitioning().partition_count() == 1 { + let stream = plan.execute(0, context)?; + let batches: Vec<RecordBatch> = stream.try_collect().await?; + return Ok(vec![batches]); + } + + let streams = execute_stream_partitioned(plan, context)?; + + let mut join_set = JoinSet::new(); + // Execute the plan and collect the results into batches. + streams.into_iter().enumerate().for_each(|(idx, stream)| { + join_set.spawn(async move { + let result: Result<Vec<RecordBatch>> = stream.try_collect().await; + (idx, result) + }); + }); + + let mut batches = vec![]; + // Note that currently this doesn't identify the thread that panicked + // + // TODO: Replace with [join_next_with_id](https://docs.rs/tokio/latest/tokio/task/struct.JoinSet.html#method.join_next_with_id + // once it is stable + while let Some(result) = join_set.join_next().await { + match result { + Ok((idx, res)) => batches.push((idx, res?)), + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } + } + } + } + + batches.sort_by_key(|(idx, _)| *idx); + let batches = batches.into_iter().map(|(_, batch)| batch).collect(); + + Ok(batches) +} + +/// Execute the [ExecutionPlan] and return a vec with one stream per output +/// partition +/// +/// # Aborting Execution +/// +/// Dropping the stream will abort the execution of the query, and free up +/// any allocated resources +#[expect( + clippy::needless_pass_by_value, + reason = "Public API that historically takes owned Arcs" +)] +pub fn execute_stream_partitioned( + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, +) -> Result<Vec<SendableRecordBatchStream>> { + let num_partitions = plan.output_partitioning().partition_count(); + let mut streams = Vec::with_capacity(num_partitions); + for i in 0..num_partitions { + streams.push(plan.execute(i, Arc::clone(&context))?); + } + Ok(streams) +} + +/// Executes an input stream and ensures that the resulting stream adheres to +/// the `not null` constraints specified in the `sink_schema`. +/// +/// # Arguments +/// +/// * `input` - An execution plan +/// * `sink_schema` - The schema to be applied to the output stream +/// * `partition` - The partition index to be executed +/// * `context` - The task context +/// +/// # Returns +/// +/// * `Result<SendableRecordBatchStream>` - A stream of `RecordBatch`es if successful +/// +/// This function first executes the given input plan for the specified partition +/// and context. It then checks if there are any columns in the input that might +/// violate the `not null` constraints specified in the `sink_schema`. If there are +/// such columns, it wraps the resulting stream to enforce the `not null` constraints +/// by invoking the [`check_not_null_constraints`] function on each batch of the stream. +#[expect( + clippy::needless_pass_by_value, + reason = "Public API that historically takes owned Arcs" +)] +pub fn execute_input_stream( + input: Arc<dyn ExecutionPlan>, + sink_schema: SchemaRef, + partition: usize, + context: Arc<TaskContext>, +) -> Result<SendableRecordBatchStream> { + let input_stream = input.execute(partition, context)?; + + debug_assert_eq!(sink_schema.fields().len(), input.schema().fields().len()); + + // Find input columns that may violate the not null constraint. + let risky_columns: Vec<_> = sink_schema + .fields() + .iter() + .zip(input.schema().fields().iter()) + .enumerate() + .filter_map(|(idx, (sink_field, input_field))| { + (!sink_field.is_nullable() && input_field.is_nullable()).then_some(idx) + }) + .collect(); + + if risky_columns.is_empty() { + Ok(input_stream) + } else { + // Check not null constraint on the input stream + Ok(Box::pin(RecordBatchStreamAdapter::new( + sink_schema, + input_stream + .map(move |batch| check_not_null_constraints(batch?, &risky_columns)), + ))) + } +} + +/// Checks a `RecordBatch` for `not null` constraints on specified columns. +/// +/// # Arguments +/// +/// * `batch` - The `RecordBatch` to be checked +/// * `column_indices` - A vector of column indices that should be checked for +/// `not null` constraints. +/// +/// # Returns +/// +/// * `Result<RecordBatch>` - The original `RecordBatch` if all constraints are met +/// +/// This function iterates over the specified column indices and ensures that none +/// of the columns contain null values. If any column contains null values, an error +/// is returned. +pub fn check_not_null_constraints( + batch: RecordBatch, + column_indices: &Vec<usize>, +) -> Result<RecordBatch> { + for &index in column_indices { + if batch.num_columns() <= index { + return exec_err!( + "Invalid batch column count {} expected > {}", + batch.num_columns(), + index + ); + } + + if batch + .column(index) + .logical_nulls() + .map(|nulls| nulls.null_count()) + .unwrap_or_default() + > 0 + { + return exec_err!( + "Invalid batch column at '{}' has null but schema specifies non-nullable", + index + ); + } + } + + Ok(batch) +} + +/// Make plan ready to be re-executed returning its clone with state reset for all nodes. +/// +/// Some plans will change their internal states after execution, making them unable to be executed again. +/// This function uses [`ExecutionPlan::reset_state`] to reset any internal state within the plan. +/// +/// An example is `CrossJoinExec`, which loads the left table into memory and stores it in the plan. +/// However, if the data of the left table is derived from the work table, it will become outdated +/// as the work table changes. When the next iteration executes this plan again, we must clear the left table. +/// +/// # Limitations +/// +/// While this function enables plan reuse, it does not allow the same plan to be executed if it (OR): +/// +/// * uses dynamic filters, +/// * represents a recursive query. +/// +pub fn reset_plan_states(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> { + plan.transform_up(|plan| { + let new_plan = Arc::clone(&plan).reset_state()?; + Ok(Transformed::yes(new_plan)) + }) + .data() +} + +/// Check if the `plan` children has the same properties as passed `children`. +/// In this case plan can avoid self properties re-computation when its children +/// replace is requested. +/// The size of `children` must be equal to the size of `ExecutionPlan::children()`. +pub fn has_same_children_properties( + plan: &dyn ExecutionPlan, + children: &[Arc<dyn ExecutionPlan>], +) -> Result<bool> { + let old_children = plan.children(); + assert_eq_or_internal_err!( + children.len(), + old_children.len(), + "Wrong number of children" + ); + for (lhs, rhs) in old_children.iter().zip(children.iter()) { + if !Arc::ptr_eq(lhs.properties(), rhs.properties()) { + return Ok(false); + } + } + Ok(true) +} + +/// Helper macro to avoid properties re-computation if passed children properties +/// the same as plan already has. Could be used to implement fast-path for method +/// [`ExecutionPlan::with_new_children`]. +/// +/// New call sites should route through [`replace_children_if_necessary`], +/// which applies this check together with the child-pointer short-circuit +/// (see [`replace_children_if_necessary`] for the layered policy). This +/// macro remains for direct-caller sites that have not been migrated yet. +#[macro_export] +macro_rules! check_if_same_properties { + ($plan: expr, $children: expr) => { + if $crate::execution_plan::has_same_children_properties( + $plan.as_ref(), + &$children, + )? { + return ::std::sync::Arc::clone(&$plan) + .with_new_children_and_same_properties($children); + } + }; +} + +/// Helper macro to validate that replacement children match a plan's existing +/// child count. +/// +/// This is useful for [`ExecutionPlan::replace_children`] implementations that +/// need to preserve the same child-count validation behavior. +#[macro_export] +macro_rules! validate_child_count { + ($plan: expr, $children: expr) => { + datafusion_common::assert_eq_or_internal_err!( + $children.len(), + $plan.children().len(), + "Wrong number of children" + ); + }; +} + +/// Utility function yielding a string representation of the given [`ExecutionPlan`]. +pub fn get_plan_string(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> { + let formatted = displayable(plan.as_ref()).indent(true).to_string(); + let actual: Vec<&str> = formatted.trim().lines().collect(); + actual.iter().map(|elem| (*elem).to_string()).collect() +} + +/// Indicates the effect an execution plan operator will have on the cardinality +/// of its input stream +pub enum CardinalityEffect { + /// Unknown effect. This is the default + Unknown, + /// The operator is guaranteed to produce exactly one row for + /// each input row + Equal, + /// The operator may produce fewer output rows than it receives input rows + LowerEqual, + /// The operator may produce more output rows than it receives input rows + GreaterEqual, +} + +/// Can be used in contexts where properties have not yet been initialized properly. +pub(crate) fn stub_properties() -> Arc<PlanProperties> { + static STUB_PROPERTIES: LazyLock<Arc<PlanProperties>> = LazyLock::new(|| { + Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )) + }); + + Arc::clone(&STUB_PROPERTIES) +} + +#[cfg(test)] +mod tests { + + use super::*; + use crate::buffer::BufferExec; + use crate::test::exec::MockExec; + use crate::{DisplayAs, DisplayFormatType, ExecutionPlan}; + + use arrow::array::{DictionaryArray, Int32Array, NullArray, RunArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; + + #[derive(Debug)] + pub struct EmptyExec { + dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>, + } + + impl EmptyExec { + pub fn new(_schema: SchemaRef) -> Self { + Self { + dynamic_expressions: vec![], + } + } + + fn with_dynamic_expressions( + mut self, + dynamic_expressions: Vec<Arc<dyn PhysicalExpr>>, + ) -> Self { + self.dynamic_expressions = dynamic_expressions; + self + } + } + + impl DisplayAs for EmptyExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for EmptyExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.dynamic_expressions.iter().map(Arc::clone).collect() + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + unimplemented!() + } + } + + #[test] + fn test_dynamic_expression_invariants() -> Result<()> { + let schema = Arc::new(Schema::empty()); + let dynamic: Arc<dyn PhysicalExpr> = + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let valid = EmptyExec::new(Arc::clone(&schema)) + .with_dynamic_expressions(vec![Arc::clone(&dynamic)]); + check_default_invariants(&valid, InvariantLevel::Always)?; + + let missing_id = + EmptyExec::new(Arc::clone(&schema)).with_dynamic_expressions(vec![lit(true)]); + let error = check_default_invariants(&missing_id, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.contains("without an expression ID"), "{error}"); + + let duplicate = EmptyExec::new(schema) + .with_dynamic_expressions(vec![Arc::clone(&dynamic), dynamic]); + let error = check_default_invariants(&duplicate, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.contains("duplicate expression ID"), "{error}"); + + Ok(()) + } + + #[derive(Debug)] + pub struct RenamedEmptyExec; + + impl RenamedEmptyExec { + pub fn new(_schema: SchemaRef) -> Self { + Self + } + } + + impl DisplayAs for RenamedEmptyExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for RenamedEmptyExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn static_name() -> &'static str + where + Self: Sized, + { + "MyRenamedEmptyExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + unimplemented!() + } + } + + #[derive(Debug)] + struct DowncastDelegatingExec(Arc<dyn ExecutionPlan>); + + impl DisplayAs for DowncastDelegatingExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for DowncastDelegatingExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + self.0.apply_expressions(f) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + Some(self.0.as_ref()) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + + fn partition_statistics( + &self, + _partition: Option<usize>, + ) -> Result<Arc<Statistics>> { + unimplemented!() + } + } + /// Test leaf plan with a real [`PlanProperties`] cache. Different instances + /// can share the same cache Arc by cloning `cache`. + #[derive(Debug, Clone)] + struct WithChildrenTestLeaf { + cache: Arc<PlanProperties>, + } + + impl WithChildrenTestLeaf { + fn new(cache: Arc<PlanProperties>) -> Self { + Self { cache } + } + } + + impl DisplayAs for WithChildrenTestLeaf { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestLeaf { + fn name(&self) -> &'static str { + "WithChildrenTestLeaf" + } + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + } + + /// Test unary plan that counts which of `with_new_children` (full + /// recompute) vs `with_new_children_and_same_properties` (fast path) is + /// taken. + #[derive(Debug, Clone)] + struct WithChildrenTestParent { + input: Arc<dyn ExecutionPlan>, + cache: Arc<PlanProperties>, + recompute_calls: Arc<std::sync::atomic::AtomicUsize>, + fast_path_calls: Arc<std::sync::atomic::AtomicUsize>, + } + + impl WithChildrenTestParent { + fn new(input: Arc<dyn ExecutionPlan>) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + input, + cache, + recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + fast_path_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + impl DisplayAs for WithChildrenTestParent { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestParent { + fn name(&self) -> &'static str { + "WithChildrenTestParent" + } + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.fast_path_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + // Full recompute: allocate a fresh `PlanProperties` Arc so this + // path is observable via `Arc::ptr_eq` on properties. + let new_input = children.swap_remove(0); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input: new_input, + cache, + recompute_calls: Arc::clone(&self.recompute_calls), + fast_path_calls: Arc::clone(&self.fast_path_calls), + })) + } + } + } + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + } + + /// Test unary plan that does **not** override + /// `with_new_children_and_same_properties`. Used to verify the default + /// trait fallback still routes through `with_new_children` (which is + /// the semantics-preserving path for downstream / external + /// `ExecutionPlan` implementations that haven't opted into the + /// fast path yet). + #[derive(Debug, Clone)] + struct WithChildrenTestParentDefault { + input: Arc<dyn ExecutionPlan>, + cache: Arc<PlanProperties>, + recompute_calls: Arc<std::sync::atomic::AtomicUsize>, + } + + impl WithChildrenTestParentDefault { + fn new(input: Arc<dyn ExecutionPlan>) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + input, + cache, + recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + impl DisplayAs for WithChildrenTestParentDefault { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestParentDefault { + fn name(&self) -> &'static str { + "WithChildrenTestParentDefault" + } + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let new_input = children.swap_remove(0); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input: new_input, + cache, + recompute_calls: Arc::clone(&self.recompute_calls), + })) + } + // Intentionally does **not** override + // `with_new_children_and_same_properties` — relies on the trait + // default that falls back to `with_new_children`. + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + } + + /// Cover the three short-circuit layers of + /// [`replace_children_if_necessary`]. + #[test] + fn test_replace_children_if_necessary_layers() -> Result<()> { + use std::sync::atomic::Ordering; + + // Two leaves that share the same `PlanProperties` Arc but sit behind + // distinct `Arc<dyn ExecutionPlan>` pointers. + let leaf_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_a: Arc<dyn ExecutionPlan> = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + let leaf_b: Arc<dyn ExecutionPlan> = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + // A third leaf with a *different* `PlanProperties` Arc — for layer 3. + let leaf_c_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_c: Arc<dyn ExecutionPlan> = + Arc::new(WithChildrenTestLeaf::new(leaf_c_props)); + + let parent = Arc::new(WithChildrenTestParent::new(Arc::clone(&leaf_a))); + let parent_dyn: Arc<dyn ExecutionPlan> = Arc::clone(&parent) as _; + let orig_props = Arc::clone(parent.properties()); + + // Layer 1: same child pointer → returns the original plan Arc verbatim. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_a)], + )?; + assert!(Arc::ptr_eq(&out, &parent_dyn)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 0); + + // Layer 2: distinct child Arc, but children share the same + // `PlanProperties` Arc → fast path, parent's `PlanProperties` cache + // Arc is reused (not reallocated). + assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); + assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_b)], + )?; + assert!(Arc::ptr_eq(out.properties(), &orig_props)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1); + + // Layer 3: child's `PlanProperties` Arc differs → full recompute. + assert!(!Arc::ptr_eq(leaf_a.properties(), leaf_c.properties())); + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_c)], + )?; + assert!(!Arc::ptr_eq(out.properties(), &orig_props)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1); + + Ok(()) + } + + /// A plan that does not override `with_new_children_and_same_properties` + /// (per @kosiew's review on #23332) must still be routed through + /// `with_new_children` when the helper hits the "same properties" + /// branch. The default trait implementation forwards to + /// `with_new_children`, so downstream / external `ExecutionPlan` + /// implementations keep the semantics-preserving path. + #[test] + fn test_replace_children_if_necessary_default_fallback() -> Result<()> { + use std::sync::atomic::Ordering; + + let leaf_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_a: Arc<dyn ExecutionPlan> = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + let leaf_b: Arc<dyn ExecutionPlan> = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); + assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); + + let parent = Arc::new(WithChildrenTestParentDefault::new(Arc::clone(&leaf_a))); + let parent_dyn: Arc<dyn ExecutionPlan> = Arc::clone(&parent) as _; + + // Using the same child means we return the original plan Arc verbatim, so even when + // the `replace_children` `ChildrenPropertiesMode::Keep` path is not defined, + // we do not recompute. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_a)], + )?; + assert!(Arc::ptr_eq(&out, &parent_dyn)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + + // Using a distinct child but the same `PlanProperties` Arc means the helper + // attempts to enter the Keep branch. If it does not exist, we fall back + // to recomputation. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_b)], + )?; + // `with_new_children` was invoked exactly once via the default. + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1); + // The returned plan has a freshly-recomputed `PlanProperties` Arc, + // so it differs from the parent's original cache. This confirms + // the fallback ran and did not short-circuit. + assert!(!Arc::ptr_eq(out.properties(), parent.properties())); + + Ok(()) + } + + /// A test node that holds a fixed list of expressions, used to test + /// `apply_expressions` behavior. + #[derive(Debug)] + struct MultiExprExec { + exprs: Vec<Arc<dyn PhysicalExpr>>, + children: Vec<Arc<dyn ExecutionPlan>>, + } + + impl DisplayAs for MultiExprExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for MultiExprExec { + fn name(&self) -> &'static str { + "MultiExprExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + self.children.iter().collect() + } + + fn with_new_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + apply_expression_roots(&self.exprs, f) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!() + } + + fn partition_statistics( + &self, + _partition: Option<usize>, + ) -> Result<Arc<Statistics>> { + unimplemented!() + } + } + + /// Returns a simple literal `Arc<dyn PhysicalExpr>` for use in tests. + fn lit_expr(val: i64) -> Arc<dyn PhysicalExpr> { + use datafusion_physical_expr::expressions::Literal; + Arc::new(Literal::new(datafusion_common::ScalarValue::Int64(Some( + val, + )))) + } + + /// `apply_expressions` visits all expressions when `f` always returns `Continue`. + #[test] + fn test_apply_expressions_continue_visits_all() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited, 3); + Ok(()) + } + + #[test] + fn test_apply_expressions_stop_halts_early() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + let tnr = plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Stop) + })?; + // Only the first expression is visited; the rest are skipped. + assert_eq!(visited, 1); + assert_eq!(tnr, TreeNodeRecursion::Stop); + Ok(()) + } + + #[test] + fn test_apply_expressions_jump_visits_next_root() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + let tnr = plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Jump) + })?; + assert_eq!(visited, 3); + assert_eq!(tnr, TreeNodeRecursion::Continue); + Ok(()) + } + + #[test] + fn test_apply_expressions_does_not_recurse() -> Result<()> { + use datafusion_physical_expr::expressions::NegativeExpr; + + let child: Arc<dyn ExecutionPlan> = Arc::new(MultiExprExec { + exprs: vec![lit_expr(2)], + children: vec![], + }); + let nested: Arc<dyn PhysicalExpr> = Arc::new(NegativeExpr::new(lit_expr(1))); + let plan = MultiExprExec { + exprs: vec![nested], + children: vec![child], + }; + + let mut visited = 0; + plan.apply_expressions(&mut |expr| { + visited += 1; + assert!(expr.is::<NegativeExpr>()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited, 1); + Ok(()) + } + + #[test] + fn test_apply_expressions_callback_can_retain_arc() -> Result<()> { + let expected = lit_expr(1); + let plan = MultiExprExec { + exprs: vec![Arc::clone(&expected)], + children: vec![], + }; + let mut retained = None; + plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(plan); + + assert!(Arc::ptr_eq( + &expected, + retained + .as_ref() + .expect("callback should retain expression") + )); + Ok(()) + } + + #[test] + fn test_execution_plan_name() { + let schema1 = Arc::new(Schema::empty()); + let default_name_exec = EmptyExec::new(schema1); + assert_eq!(default_name_exec.name(), "EmptyExec"); + + let schema2 = Arc::new(Schema::empty()); + let renamed_exec = RenamedEmptyExec::new(schema2); + assert_eq!(renamed_exec.name(), "MyRenamedEmptyExec"); + assert_eq!(RenamedEmptyExec::static_name(), "MyRenamedEmptyExec"); + } + + #[test] + fn test_execution_plan_downcast_delegates_to_downcast_delegate() { + let schema = Arc::new(Schema::empty()); + let inner: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(schema)); + let wrapped: Arc<dyn ExecutionPlan> = Arc::new(DowncastDelegatingExec(inner)); + let nested: Arc<dyn ExecutionPlan> = + Arc::new(DowncastDelegatingExec(Arc::clone(&wrapped))); + + for plan in [wrapped.as_ref(), nested.as_ref()] { + assert!(!plan.is::<DowncastDelegatingExec>()); + assert!(plan.downcast_ref::<DowncastDelegatingExec>().is_none()); + assert!(plan.is::<EmptyExec>()); + assert!(plan.downcast_ref::<EmptyExec>().is_some()); + assert!(!plan.is::<RenamedEmptyExec>()); + assert!(plan.downcast_ref::<RenamedEmptyExec>().is_none()); + } + } + + /// A compilation test to ensure that the `ExecutionPlan::name()` method can + /// be called from a trait object. + /// Related ticket: https://github.com/apache/datafusion/pull/11047 + #[expect(unused)] + fn use_execution_plan_as_trait_object(plan: &dyn ExecutionPlan) { + let _ = plan.name(); + } + + #[test] + fn buffer_exec_does_not_need_data_exchange() { + let schema = Arc::new(Schema::empty()); + let input: Arc<dyn ExecutionPlan> = Arc::new(MockExec::new(vec![], schema)); + let buffer: Arc<dyn ExecutionPlan> = Arc::new(BufferExec::new(input, 1024)); + + assert!(!need_data_exchange(buffer)); + } + + #[test] + fn test_check_not_null_constraints_accept_non_null() -> Result<()> { + check_not_null_constraints( + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)]))], + )?, + &vec![0], + )?; + Ok(()) + } + + #[test] + fn test_check_not_null_constraints_reject_null() -> Result<()> { + let result = check_not_null_constraints( + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![Some(1), None, Some(3)]))], + )?, + &vec![0], + ); + assert!(result.is_err()); + assert_eq!( + result.err().unwrap().strip_backtrace(), + "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable", + ); + Ok(()) + } + + #[test] + fn test_check_not_null_constraints_with_run_end_array() -> Result<()> { + // some null value inside REE array + let run_ends = Int32Array::from(vec![1, 2, 3, 4]); + let values = Int32Array::from(vec![Some(0), None, Some(1), None]); + let run_end_array = RunArray::try_new(&run_ends, &values)?; + let result = check_not_null_constraints( + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "a", + run_end_array.data_type().to_owned(), + true, + )])), + vec![Arc::new(run_end_array)], + )?, + &vec![0], + ); + assert!(result.is_err()); + assert_eq!( + result.err().unwrap().strip_backtrace(), + "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable", + ); + Ok(()) + } + + #[test] + fn test_check_not_null_constraints_with_dictionary_array_with_null() -> Result<()> { + let values = Arc::new(Int32Array::from(vec![Some(1), None, Some(3), Some(4)])); + let keys = Int32Array::from(vec![0, 1, 2, 3]); + let dictionary = DictionaryArray::new(keys, values); + let result = check_not_null_constraints( + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "a", + dictionary.data_type().to_owned(), + true, + )])), + vec![Arc::new(dictionary)], + )?, + &vec![0], + ); + assert!(result.is_err()); + assert_eq!( + result.err().unwrap().strip_backtrace(), + "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable", + ); + Ok(()) + } + + #[test] + fn test_check_not_null_constraints_with_dictionary_masking_null() -> Result<()> { + // some null value marked out by dictionary array + let values = Arc::new(Int32Array::from(vec![ + Some(1), + None, // this null value is masked by dictionary keys + Some(3), + Some(4), + ])); + let keys = Int32Array::from(vec![0, /*1,*/ 2, 3]); + let dictionary = DictionaryArray::new(keys, values); + check_not_null_constraints( + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "a", + dictionary.data_type().to_owned(), + true, + )])), + vec![Arc::new(dictionary)], + )?, + &vec![0], + )?; + Ok(()) + } + + #[test] + fn test_check_not_null_constraints_on_null_type() -> Result<()> { + // null value of Null type + let result = check_not_null_constraints( + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Null, true)])), + vec![Arc::new(NullArray::new(3))], + )?, + &vec![0], + ); + assert!(result.is_err()); + assert_eq!( + result.err().unwrap().strip_backtrace(), + "Execution error: Invalid batch column at '0' has null but schema specifies non-nullable", + ); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/explain.rs b/native/vendor/datafusion-physical-plan/src/explain.rs new file mode 100644 index 0000000..3b31ee7 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/explain.rs
@@ -0,0 +1,403 @@ +// 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. + +//! Defines the EXPLAIN operator + +use std::sync::Arc; + +use super::{DisplayAs, PlanProperties, SendableRecordBatchStream}; +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, +}; + +use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; +use datafusion_common::display::StringifiedPlan; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; + +use log::trace; + +/// Explain execution plan operator. This operator contains the string +/// values of the various plans it has when it is created, and passes +/// them to its output. +#[derive(Debug, Clone)] +pub struct ExplainExec { + /// The schema that this exec plan node outputs + schema: SchemaRef, + /// The strings to be printed + stringified_plans: Vec<StringifiedPlan>, + /// control which plans to print + verbose: bool, + cache: Arc<PlanProperties>, +} + +impl ExplainExec { + /// Create a new ExplainExec + pub fn new( + schema: SchemaRef, + stringified_plans: Vec<StringifiedPlan>, + verbose: bool, + ) -> Self { + let cache = Self::compute_properties(Arc::clone(&schema)); + ExplainExec { + schema, + stringified_plans, + verbose, + cache: Arc::new(cache), + } + } + + /// The strings to be printed + pub fn stringified_plans(&self) -> &[StringifiedPlan] { + &self.stringified_plans + } + + /// Access to verbose + pub fn verbose(&self) -> bool { + self.verbose + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for ExplainExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ExplainExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for ExplainExec { + fn name(&self) -> &'static str { + "ExplainExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + // This is a leaf node and has no children + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start ExplainExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + assert_eq_or_internal_err!( + partition, + 0, + "ExplainExec invalid partition {partition}" + ); + let mut type_builder = + StringBuilder::with_capacity(self.stringified_plans.len(), 1024); + let mut plan_builder = + StringBuilder::with_capacity(self.stringified_plans.len(), 1024); + + let plans_to_print = self + .stringified_plans + .iter() + .filter(|s| s.should_display(self.verbose)); + + // Identify plans that are not changed + let mut prev: Option<&StringifiedPlan> = None; + + for p in plans_to_print { + type_builder.append_value(p.plan_type.to_string()); + match prev { + Some(prev) if !should_show(prev, p) => { + plan_builder.append_value("SAME TEXT AS ABOVE"); + } + Some(_) | None => { + plan_builder.append_value(&*p.plan); + } + } + prev = Some(p); + } + + let record_batch = RecordBatch::try_new( + Arc::clone(&self.schema), + vec![ + Arc::new(type_builder.finish()), + Arc::new(plan_builder.finish()), + ], + )?; + + trace!( + "Before returning RecordBatchStream in ExplainExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + futures::stream::iter(vec![Ok(record_batch)]), + ))) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Explain( + protobuf::ExplainExecNode { + schema: Some(self.schema().as_ref().try_into()?), + stringified_plans: self + .stringified_plans() + .iter() + .map(stringified_plan_to_proto) + .collect(), + verbose: self.verbose(), + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ExplainExec { + /// Reconstruct an [`ExplainExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let explain = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Explain, + "ExplainExec", + ); + let schema = explain.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ExplainExec is missing required field 'schema'" + ) + })?; + Ok(Arc::new(ExplainExec::new( + Arc::new(arrow::datatypes::Schema::try_from(schema)?), + explain + .stringified_plans + .iter() + .map(stringified_plan_from_proto) + .collect(), + explain.verbose, + ))) + } +} + +#[cfg(feature = "proto")] +fn stringified_plan_to_proto( + stringified_plan: &StringifiedPlan, +) -> datafusion_proto_models::protobuf::StringifiedPlan { + use datafusion_common::display::PlanType; + use datafusion_proto_models::datafusion_common::EmptyMessage; + use datafusion_proto_models::protobuf; + use protobuf::plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }; + + protobuf::StringifiedPlan { + plan_type: match stringified_plan.clone().plan_type { + PlanType::InitialLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialLogicalPlan(EmptyMessage {})), + }), + PlanType::AnalyzedLogicalPlan { analyzer_name } => Some(protobuf::PlanType { + plan_type_enum: Some(AnalyzedLogicalPlan( + protobuf::AnalyzedLogicalPlanType { analyzer_name }, + )), + }), + PlanType::FinalAnalyzedLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalAnalyzedLogicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedLogicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedLogicalPlan( + protobuf::OptimizedLogicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalLogicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedPhysicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedPhysicalPlan( + protobuf::OptimizedPhysicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::PhysicalPlanError => Some(protobuf::PlanType { + plan_type_enum: Some(PhysicalPlanError(EmptyMessage {})), + }), + }, + plan: stringified_plan.plan.to_string(), + } +} + +#[cfg(feature = "proto")] +fn stringified_plan_from_proto( + stringified_plan: &datafusion_proto_models::protobuf::StringifiedPlan, +) -> StringifiedPlan { + use datafusion_common::display::PlanType; + use datafusion_proto_models::protobuf::plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }; + use datafusion_proto_models::protobuf::{ + AnalyzedLogicalPlanType, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, + }; + + StringifiedPlan { + plan_type: match stringified_plan + .plan_type + .as_ref() + .and_then(|plan_type| plan_type.plan_type_enum.as_ref()) + .unwrap_or_else(|| { + panic!( + "Cannot create protobuf::StringifiedPlan from {stringified_plan:?}" + ) + }) { + InitialLogicalPlan(_) => PlanType::InitialLogicalPlan, + AnalyzedLogicalPlan(AnalyzedLogicalPlanType { analyzer_name }) => { + PlanType::AnalyzedLogicalPlan { + analyzer_name: analyzer_name.clone(), + } + } + FinalAnalyzedLogicalPlan(_) => PlanType::FinalAnalyzedLogicalPlan, + OptimizedLogicalPlan(OptimizedLogicalPlanType { optimizer_name }) => { + PlanType::OptimizedLogicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalLogicalPlan(_) => PlanType::FinalLogicalPlan, + InitialPhysicalPlan(_) => PlanType::InitialPhysicalPlan, + InitialPhysicalPlanWithStats(_) => PlanType::InitialPhysicalPlanWithStats, + InitialPhysicalPlanWithSchema(_) => PlanType::InitialPhysicalPlanWithSchema, + OptimizedPhysicalPlan(OptimizedPhysicalPlanType { optimizer_name }) => { + PlanType::OptimizedPhysicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalPhysicalPlan(_) => PlanType::FinalPhysicalPlan, + FinalPhysicalPlanWithStats(_) => PlanType::FinalPhysicalPlanWithStats, + FinalPhysicalPlanWithSchema(_) => PlanType::FinalPhysicalPlanWithSchema, + PhysicalPlanError(_) => PlanType::PhysicalPlanError, + }, + plan: Arc::new(stringified_plan.plan.clone()), + } +} + +/// If this plan should be shown, given the previous plan that was +/// displayed. +/// +/// This is meant to avoid repeating the same plan over and over again +/// in explain plans to make clear what is changing +fn should_show(previous_plan: &StringifiedPlan, this_plan: &StringifiedPlan) -> bool { + // if the plans are different, or if they would have been + // displayed in the normal explain (aka non verbose) plan + (previous_plan.plan != this_plan.plan) || this_plan.should_display(false) +}
diff --git a/native/vendor/datafusion-physical-plan/src/filter.rs b/native/vendor/datafusion-physical-plan/src/filter.rs new file mode 100644 index 0000000..5df5482 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/filter.rs
@@ -0,0 +1,3907 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::hash_map::Entry; +use std::collections::{HashMap, HashSet}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, ready}; + +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; +use itertools::Itertools; + +use super::{ + ColumnStatistics, DisplayAs, ExecutionPlanProperties, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, Statistics, +}; +use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; +use crate::common::can_project; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, PushedDown, +}; +use crate::limit::LocalLimitExec; +use crate::metrics::{MetricBuilder, MetricType}; +use crate::projection::{ + EmbeddedProjection, ProjectionExec, ProjectionExpr, make_with_child, + try_embed_projection, update_expr, +}; +use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; +use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; +use crate::{ + DisplayFormatType, ExecutionPlan, + metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RatioMetrics}, +}; + +use arrow::compute::filter_record_batch; +use arrow::datatypes::{DataType, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::cast::as_boolean_array; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, internal_err, plan_err, project_schema, +}; +use datafusion_execution::TaskContext; +use datafusion_expr::Operator; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr::expressions::{ + BinaryExpr, Column, IsNotNullExpr, Literal, lit, +}; +use datafusion_physical_expr::intervals::utils::check_support; +use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns}; +use datafusion_physical_expr::{ + AcrossPartitions, AnalysisContext, ConstExpr, ExprBoundaries, PhysicalExpr, analyze, + conjunction, split_conjunction, +}; + +use datafusion_physical_expr_common::physical_expr::fmt_sql; +use futures::stream::{Stream, StreamExt}; +use log::trace; + +const FILTER_EXEC_DEFAULT_SELECTIVITY: u8 = 20; +const FILTER_EXEC_DEFAULT_BATCH_SIZE: usize = 8192; + +/// FilterExec evaluates a boolean predicate against all input batches to determine which rows to +/// include in its output batches. +#[derive(Debug, Clone)] +pub struct FilterExec { + /// The expression to filter on. This expression must evaluate to a boolean value. + predicate: Arc<dyn PhysicalExpr>, + /// The input plan + input: Arc<dyn ExecutionPlan>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Selectivity for statistics. 0 = no rows, 100 = all rows + default_selectivity: u8, + /// Properties equivalence properties, partitioning, etc. + cache: Arc<PlanProperties>, + /// The projection indices of the columns in the output schema of join + projection: Option<ProjectionRef>, + /// Target batch size for output batches + batch_size: usize, + /// Number of rows to fetch + fetch: Option<usize>, +} + +/// Builder for [`FilterExec`] to set optional parameters +pub struct FilterExecBuilder { + predicate: Arc<dyn PhysicalExpr>, + input: Arc<dyn ExecutionPlan>, + projection: Option<ProjectionRef>, + default_selectivity: u8, + batch_size: usize, + fetch: Option<usize>, +} + +impl FilterExecBuilder { + /// Create a new builder with required parameters (predicate and input) + pub fn new(predicate: Arc<dyn PhysicalExpr>, input: Arc<dyn ExecutionPlan>) -> Self { + Self { + predicate, + input, + projection: None, + default_selectivity: FILTER_EXEC_DEFAULT_SELECTIVITY, + batch_size: FILTER_EXEC_DEFAULT_BATCH_SIZE, + fetch: None, + } + } + + /// Set the input execution plan + pub fn with_input(mut self, input: Arc<dyn ExecutionPlan>) -> Self { + self.input = input; + self + } + + /// Set the predicate expression + pub fn with_predicate(mut self, predicate: Arc<dyn PhysicalExpr>) -> Self { + self.predicate = predicate; + self + } + + /// Set the projection, composing with any existing projection. + /// + /// If a projection is already set, the new projection indices are mapped + /// through the existing projection. For example, if the current projection + /// is `[0, 2, 3]` and `apply_projection(Some(vec![0, 2]))` is called, the + /// resulting projection will be `[0, 3]` (indices 0 and 2 of `[0, 2, 3]`). + /// + /// If no projection is currently set, the new projection is used directly. + /// If `None` is passed, the projection is cleared. + pub fn apply_projection(self, projection: Option<Vec<usize>>) -> Result<Self> { + let projection = projection.map(Into::into); + self.apply_projection_by_ref(projection.as_ref()) + } + + /// The same as [`Self::apply_projection`] but takes projection shared reference. + pub fn apply_projection_by_ref( + mut self, + projection: Option<&ProjectionRef>, + ) -> Result<Self> { + // Check if the projection is valid against current output schema + can_project(&self.input.schema(), projection.map(AsRef::as_ref))?; + self.projection = combine_projections(projection, self.projection.as_ref())?; + Ok(self) + } + + /// Set the default selectivity + pub fn with_default_selectivity(mut self, default_selectivity: u8) -> Self { + self.default_selectivity = default_selectivity; + self + } + + /// Set the batch size + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = batch_size; + self + } + + /// Set the fetch limit + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + /// Build the FilterExec, computing properties once with all configured parameters + pub fn build(self) -> Result<FilterExec> { + // Validate predicate type + match self.predicate.data_type(self.input.schema().as_ref())? { + DataType::Boolean => {} + other => { + return plan_err!( + "Filter predicate must return BOOLEAN values, got {other:?}" + ); + } + } + + // Validate selectivity + if self.default_selectivity > 100 { + return plan_err!( + "Default filter selectivity value needs to be less than or equal to 100" + ); + } + + // Validate projection if provided + can_project(&self.input.schema(), self.projection.as_deref())?; + + // Compute properties once with all parameters + let cache = FilterExec::compute_properties( + &self.input, + &self.predicate, + self.default_selectivity, + self.projection.as_deref(), + )?; + + Ok(FilterExec { + predicate: self.predicate, + input: self.input, + metrics: ExecutionPlanMetricsSet::new(), + default_selectivity: self.default_selectivity, + cache: Arc::new(cache), + projection: self.projection, + batch_size: self.batch_size, + fetch: self.fetch, + }) + } +} + +impl From<&FilterExec> for FilterExecBuilder { + fn from(exec: &FilterExec) -> Self { + Self { + predicate: Arc::clone(&exec.predicate), + input: Arc::clone(&exec.input), + projection: exec.projection.clone(), + default_selectivity: exec.default_selectivity, + batch_size: exec.batch_size, + fetch: exec.fetch, + // We could cache / copy over PlanProperties + // here but that would require invalidating them in FilterExecBuilder::apply_projection, etc. + // and currently every call to this method ends up invalidating them anyway. + // If useful this can be added in the future as a non-breaking change. + } + } +} + +impl FilterExec { + /// Create a FilterExec on an input using the builder pattern + pub fn try_new( + predicate: Arc<dyn PhysicalExpr>, + input: Arc<dyn ExecutionPlan>, + ) -> Result<Self> { + FilterExecBuilder::new(predicate, input).build() + } + + /// Get a batch size + pub fn batch_size(&self) -> usize { + self.batch_size + } + + /// Set the default selectivity + pub fn with_default_selectivity( + mut self, + default_selectivity: u8, + ) -> Result<Self, DataFusionError> { + if default_selectivity > 100 { + return plan_err!( + "Default filter selectivity value needs to be less than or equal to 100" + ); + } + self.default_selectivity = default_selectivity; + Ok(self) + } + + /// Return new instance of [FilterExec] with the given projection. + /// + /// # Deprecated + /// Use [`FilterExecBuilder::apply_projection`] instead + #[deprecated( + since = "52.0.0", + note = "Use FilterExecBuilder::apply_projection instead" + )] + pub fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> { + let builder = FilterExecBuilder::from(self); + builder.apply_projection(projection)?.build() + } + + /// Set the batch size + pub fn with_batch_size(&self, batch_size: usize) -> Result<Self> { + Ok(Self { + predicate: Arc::clone(&self.predicate), + input: Arc::clone(&self.input), + metrics: self.metrics.clone(), + default_selectivity: self.default_selectivity, + cache: Arc::clone(&self.cache), + projection: self.projection.clone(), + batch_size, + fetch: self.fetch, + }) + } + + /// The expression to filter on. This expression must evaluate to a boolean value. + pub fn predicate(&self) -> &Arc<dyn PhysicalExpr> { + &self.predicate + } + + /// The input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// The default selectivity + pub fn default_selectivity(&self) -> u8 { + self.default_selectivity + } + + /// Projection + pub fn projection(&self) -> &Option<ProjectionRef> { + &self.projection + } + + /// Calculates `Statistics` for `FilterExec` by applying the filter's + /// selectivity (default, or estimated from interval analysis) to the input + /// statistics. + /// + /// The estimated output row count is used to keep the per-column statistics + /// consistent with it: + /// - null and distinct counts are capped at the estimated row count; + /// - byte sizes (per column and total) are scaled by the selectivity, and + /// are an exact zero when the row count is an exact zero; + /// - a column constrained to a single value (`col = literal`, or an + /// interval that collapses to one point) gets a distinct count of 1; + /// - a column in a null-rejecting conjunct gets a null count of 0. + /// + /// When interval analysis applies, min/max are also tightened to the + /// surviving value range. + /// + /// A contradictory predicate (e.g. `a = 1 AND a = 2`) yields zero rows and + /// empty-column statistics. + pub(crate) fn statistics_helper( + schema: &SchemaRef, + input_stats: Statistics, + predicate: &Arc<dyn PhysicalExpr>, + default_selectivity: u8, + ) -> Result<Statistics> { + let (eq_columns, is_infeasible) = collect_equality_columns(predicate); + + let input_num_rows = input_stats.num_rows; + let input_total_byte_size = input_stats.total_byte_size; + + let (selectivity, num_rows, column_statistics) = if is_infeasible { + // Contradictory predicate: no rows survive. Row-bounded counts are + // zero; value statistics are undefined on an empty column. + let mut cs = input_stats.to_inexact().column_statistics; + for col_stat in &mut cs { + col_stat.distinct_count = Precision::Exact(0); + col_stat.null_count = Precision::Exact(0); + col_stat.min_value = Precision::Absent; + col_stat.max_value = Precision::Absent; + col_stat.sum_value = Precision::Absent; + col_stat.byte_size = Precision::Exact(0); + } + (0.0, Precision::Exact(0), cs) + } else { + let null_rejecting_columns = collect_null_rejecting_columns(predicate); + + if check_support(predicate, schema) { + let input_analysis_ctx = AnalysisContext::try_from_statistics( + schema, + &input_stats.column_statistics, + )?; + let analysis_ctx = analyze(predicate, input_analysis_ctx, schema)?; + let selectivity = analysis_ctx.selectivity.unwrap_or(1.0); + let filtered_num_rows = + input_num_rows.with_estimated_selectivity(selectivity); + let cs = collect_new_statistics( + schema, + &input_stats.column_statistics, + analysis_ctx.boundaries, + selectivity, + &null_rejecting_columns, + filtered_num_rows, + ); + (selectivity, filtered_num_rows, cs) + } else { + // Without interval boundaries, use the default selectivity and + // apply the row-count constraints that still follow from the + // filter predicate. + let selectivity = default_selectivity as f64 / 100.0; + let filtered_num_rows = + input_num_rows.with_estimated_selectivity(selectivity); + let mut cs = input_stats.to_inexact().column_statistics; + for (idx, col_stat) in cs.iter_mut().enumerate() { + col_stat.byte_size = scale_byte_size_at_rows( + col_stat.byte_size, + selectivity, + filtered_num_rows, + ); + col_stat.null_count = if null_rejecting_columns.contains(&idx) { + Precision::Exact(0) + } else { + cap_at_rows(col_stat.null_count, filtered_num_rows) + }; + col_stat.distinct_count = if eq_columns.contains(&idx) { + distinct_count_for_singleton_domain(filtered_num_rows) + } else { + cap_at_rows(col_stat.distinct_count, filtered_num_rows) + }; + } + (selectivity, filtered_num_rows, cs) + } + }; + + let total_byte_size = + scale_byte_size_at_rows(input_total_byte_size, selectivity, num_rows); + + Ok(Statistics { + num_rows, + total_byte_size, + column_statistics, + }) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + predicate: &Arc<dyn PhysicalExpr>, + default_selectivity: u8, + projection: Option<&[usize]>, + ) -> Result<PlanProperties> { + // Combine the equal predicates with the input equivalence properties + // to construct the equivalence properties: + let schema = input.schema(); + let stats = Self::statistics_helper( + &schema, + Arc::unwrap_or_clone( + StatisticsContext::new() + .compute(input.as_ref(), &StatisticsArgs::new())?, + ), + predicate, + default_selectivity, + )?; + let mut eq_properties = input.equivalence_properties().clone(); + let (equal_pairs, _) = collect_columns_from_predicate_inner(predicate); + for (lhs, rhs) in equal_pairs { + eq_properties.add_equal_conditions(Arc::clone(lhs), Arc::clone(rhs))? + } + // Add the columns that have only one viable value (singleton) after + // filtering to constants. + let constants = collect_columns(predicate) + .into_iter() + .filter(|column| stats.column_statistics[column.index()].is_singleton()) + .map(|column| { + let value = stats.column_statistics[column.index()] + .min_value + .get_value(); + let expr = Arc::new(column) as _; + ConstExpr::new(expr, AcrossPartitions::Uniform(value.cloned())) + }); + // This is for statistics + eq_properties.add_constants(constants)?; + // This is for logical constant (for example: a = '1', then a could be marked as a constant) + // to do: how to deal with multiple situation to represent = (for example c1 between 0 and 0) + eq_properties.add_constants(ConstExpr::collect_predicate_constants( + input.equivalence_properties(), + predicate, + ))?; + + let mut output_partitioning = input.output_partitioning().clone(); + // If contains projection, update the PlanProperties. + if let Some(projection) = projection { + let schema = eq_properties.schema(); + let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; + let out_schema = project_schema(schema, Some(&projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, out_schema); + } + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + input.pipeline_behavior(), + input.boundedness(), + )) + } +} + +impl DisplayAs for FilterExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let display_projections = if let Some(projection) = + self.projection.as_ref() + { + format!( + ", projection=[{}]", + projection + .iter() + .map(|index| format!( + "{}@{}", + self.input.schema().fields().get(*index).unwrap().name(), + index + )) + .collect::<Vec<_>>() + .join(", ") + ) + } else { + "".to_string() + }; + let fetch = self + .fetch + .map_or_else(|| "".to_string(), |f| format!(", fetch={f}")); + write!( + f, + "FilterExec: {}{}{}", + self.predicate, display_projections, fetch + ) + } + DisplayFormatType::TreeRender => { + if let Some(fetch) = self.fetch { + writeln!(f, "fetch={fetch}")?; + } + write!(f, "predicate={}", fmt_sql(self.predicate.as_ref())) + } + } + } +} + +impl ExecutionPlan for FilterExec { + fn name(&self) -> &'static str { + "FilterExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + crate::apply_expression_roots([&self.predicate], f) + } + + fn maintains_input_order(&self) -> Vec<bool> { + // Tell optimizer this operator doesn't reorder its input + vec![true] + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start FilterExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + let metrics = FilterExecMetrics::new(&self.metrics, partition); + Ok(Box::pin(FilterExecStream { + schema: self.schema(), + predicate: Arc::clone(&self.predicate), + input: self.input.execute(partition, context)?, + metrics, + projection: self.projection.clone(), + batch_coalescer: LimitedBatchCoalescer::new( + self.schema(), + self.batch_size, + self.fetch, + ), + })) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + /// The output statistics of a filtering operation can be estimated if the + /// predicate's selectivity value can be determined for the incoming data. + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let input_stats = input_stats[0].as_ref().clone(); + let stats = Self::statistics_helper( + &self.input.schema(), + input_stats, + self.predicate(), + self.default_selectivity, + )?; + Ok(Arc::new(stats.project(self.projection.as_ref()))) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + /// Tries to swap `projection` with its input (`filter`). If possible, performs + /// the swap and returns [`FilterExec`] as the top plan. Otherwise, returns `None`. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection does not narrow the schema, we should not try to push it down: + if projection.expr().len() < projection.input().schema().fields().len() { + // Each column in the predicate expression must exist after the projection. + if let Some(new_predicate) = + update_expr(self.predicate(), projection.expr(), false)? + { + return FilterExecBuilder::from(self) + .with_input(make_with_child(projection, self.input())?) + .with_predicate(new_predicate) + // The original FilterExec projection referenced columns from its old + // input. After the swap the new input is the ProjectionExec which + // already handles column selection, so clear the projection here. + .apply_projection(None)? + .build() + .map(|e| Some(Arc::new(e) as _)); + } + } + try_embed_projection(projection, self) + } + + fn gather_filters_for_pushdown( + &self, + phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + if phase != FilterPushdownPhase::Pre { + let child = + ChildFilterDescription::from_child(&parent_filters, self.input())?; + return Ok(FilterDescription::new().with_child(child)); + } + + let child = ChildFilterDescription::from_child(&parent_filters, self.input())? + .with_self_filters( + split_conjunction(&self.predicate) + .into_iter() + .cloned() + .collect(), + ); + + Ok(FilterDescription::new().with_child(child)) + } + + fn handle_child_pushdown_result( + &self, + phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + if phase != FilterPushdownPhase::Pre { + return Ok(FilterPushdownPropagation::if_all(child_pushdown_result)); + } + // We absorb any parent filters that were not handled by our children + let mut unsupported_parent_filters: Vec<Arc<dyn PhysicalExpr>> = + child_pushdown_result + .parent_filters + .iter() + .filter_map(|f| { + matches!(f.all(), PushedDown::No).then_some(Arc::clone(&f.filter)) + }) + .collect(); + + // If this FilterExec has a projection, the unsupported parent filters + // are in the output schema (after projection) coordinates. We need to + // remap them to the input schema coordinates before combining with self filters. + if self.projection.is_some() { + let input_schema = self.input().schema(); + unsupported_parent_filters = unsupported_parent_filters + .into_iter() + .map(|expr| reassign_expr_columns(expr, &input_schema)) + .collect::<Result<Vec<_>>>()?; + } + + let unsupported_self_filters = child_pushdown_result + .self_filters + .first() + .expect("we have exactly one child") + .iter() + .filter_map(|f| match f.discriminant { + PushedDown::Yes => None, + PushedDown::No => Some(&f.predicate), + }) + .cloned(); + + let unhandled_filters = unsupported_parent_filters + .into_iter() + .chain(unsupported_self_filters) + .collect_vec(); + + // If we have unhandled filters, we need to create a new FilterExec + let filter_input = Arc::clone(self.input()); + let new_predicate = conjunction(unhandled_filters); + let updated_node = if new_predicate.eq(&lit(true)) { + // FilterExec is no longer needed, but we may need to leave a projection in place. + // If this FilterExec had a fetch limit, propagate it to the child. + // When the child also has a fetch, use the minimum of both to preserve + // the tighter constraint. + let filter_input = if let Some(outer_fetch) = self.fetch { + let effective_fetch = match filter_input.fetch() { + Some(inner_fetch) => outer_fetch.min(inner_fetch), + None => outer_fetch, + }; + match filter_input.with_fetch(Some(effective_fetch)) { + Some(node) => node, + None => Arc::new(LocalLimitExec::new(filter_input, effective_fetch)), + } + } else { + filter_input + }; + match self.projection().as_ref() { + Some(projection_indices) => { + let filter_child_schema = filter_input.schema(); + let proj_exprs = projection_indices + .iter() + .map(|p| { + let field = filter_child_schema.field(*p).clone(); + ProjectionExpr { + expr: Arc::new(Column::new(field.name(), *p)) + as Arc<dyn PhysicalExpr>, + alias: field.name().to_string(), + } + }) + .collect::<Vec<_>>(); + Some(Arc::new(ProjectionExec::try_new(proj_exprs, filter_input)?) + as Arc<dyn ExecutionPlan>) + } + None => { + // No projection needed, just return the input + Some(filter_input) + } + } + } else if new_predicate.eq(&self.predicate) { + // The new predicate is the same as our current predicate + None + } else { + // Create a new FilterExec with the new predicate, preserving the projection + let new = FilterExec { + predicate: Arc::clone(&new_predicate), + input: Arc::clone(&filter_input), + metrics: self.metrics.clone(), + default_selectivity: self.default_selectivity, + cache: Arc::new(Self::compute_properties( + &filter_input, + &new_predicate, + self.default_selectivity, + self.projection.as_deref(), + )?), + projection: self.projection.clone(), + batch_size: self.batch_size, + fetch: self.fetch, + }; + Some(Arc::new(new) as _) + }; + + Ok(FilterPushdownPropagation { + filters: vec![PushedDown::Yes; child_pushdown_result.parent_filters.len()], + updated_node, + }) + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn with_fetch(&self, fetch: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + Some(Arc::new(Self { + predicate: Arc::clone(&self.predicate), + input: Arc::clone(&self.input), + metrics: self.metrics.clone(), + default_selectivity: self.default_selectivity, + cache: Arc::clone(&self.cache), + projection: self.projection.clone(), + batch_size: self.batch_size, + fetch, + })) + } + + fn with_preserve_order( + &self, + preserve_order: bool, + ) -> Option<Arc<dyn ExecutionPlan>> { + self.input + .with_preserve_order(preserve_order) + .and_then(|new_input| { + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) + .ok() + }) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expr(self.predicate())?; + // Preserve the exact wire format: `None` (full projection) is serialized + // as the identity projection `[0, 1, ..., num_fields - 1]` so that it is + // distinguishable from an explicit projection on decode. + let projection = if let Some(v) = self.projection() { + v.iter().map(|x| *x as u32).collect() + } else { + (0..self.input().schema().fields().len()) + .map(|i| i as u32) + .collect() + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Filter(Box::new( + protobuf::FilterExecNode { + input: Some(Box::new(input)), + expr: Some(expr), + default_filter_selectivity: self.default_selectivity() as u32, + projection, + batch_size: self.batch_size() as u32, + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl FilterExec { + /// Reconstruct a [`FilterExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one signature. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let filter = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Filter, + "FilterExec", + ); + let input = + ctx.decode_required_child(filter.input.as_deref(), "FilterExec", "input")?; + let predicate = ctx.decode_required_expr( + filter.expr.as_ref(), + input.schema().as_ref(), + "FilterExec", + "expr", + )?; + let filter_selectivity = filter.default_filter_selectivity.try_into(); + + // `None` is encoded as the full identity projection. Reconstruct it only + // when all input columns are present in order, leaving an empty list as + // `Some(vec![])`. + let num_fields = input.schema().fields().len(); + let mut is_full_projection = filter.projection.len() == num_fields; + let mut projection_vec: Vec<usize> = Vec::with_capacity(filter.projection.len()); + for (i, idx) in filter.projection.iter().enumerate() { + let idx = *idx as usize; + is_full_projection &= idx == i; + projection_vec.push(idx); + } + let projection = if is_full_projection { + None + } else { + Some(projection_vec) + }; + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(projection)? + .with_batch_size(filter.batch_size as usize) + .with_fetch(filter.fetch.map(|f| f as usize)) + .build()?; + match filter_selectivity { + Ok(filter_selectivity) => Ok(Arc::new( + filter.with_default_selectivity(filter_selectivity)?, + )), + Err(_) => Err(datafusion_common::internal_datafusion_err!( + "filter_selectivity in PhysicalPlanNode is invalid" + )), + } + } +} + +impl EmbeddedProjection for FilterExec { + fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> { + FilterExecBuilder::from(self) + .apply_projection(projection)? + .build() + } +} + +/// Collects column equality information from `col = literal` predicates in a +/// conjunction. +/// +/// Returns `(eq_columns, is_infeasible)`: +/// - `eq_columns`: set of column indices constrained to a single literal value. +/// - `is_infeasible`: `true` when the same column is equated to two different +/// non-null literals (e.g. `name = 'alice' AND name = 'bob'`), which is +/// always unsatisfiable. +/// +/// Only AND conjunctions are traversed; OR is intentionally skipped +/// since `a = 1 OR a = 2` does not pin NDV to 1. +fn collect_equality_columns(predicate: &Arc<dyn PhysicalExpr>) -> (HashSet<usize>, bool) { + let mut eq_values: HashMap<usize, ScalarValue> = HashMap::new(); + let mut infeasible = false; + + for expr in split_conjunction(predicate) { + let Some(binary) = expr.downcast_ref::<BinaryExpr>() else { + continue; + }; + if *binary.op() != Operator::Eq { + continue; + } + let left = binary.left(); + let right = binary.right(); + let pair = if let Some(col) = left.downcast_ref::<Column>() + && let Some(lit) = right.downcast_ref::<Literal>() + && !lit.value().is_null() + { + Some((col.index(), lit.value().clone())) + } else if let Some(col) = right.downcast_ref::<Column>() + && let Some(lit) = left.downcast_ref::<Literal>() + && !lit.value().is_null() + { + Some((col.index(), lit.value().clone())) + } else { + None + }; + + if let Some((idx, value)) = pair { + match eq_values.entry(idx) { + Entry::Occupied(prev) => { + if *prev.get() != value { + infeasible = true; + break; + } + } + Entry::Vacant(slot) => { + slot.insert(value); + } + } + } + } + + (eq_values.into_keys().collect(), infeasible) +} + +/// Collects columns that cannot be NULL in any surviving row. +/// +/// A filter keeps only rows where the predicate is TRUE, so a column is +/// null-rejecting if some top-level AND conjunct evaluates to NULL or FALSE +/// whenever that column is NULL. Two such conjuncts are recognized: +/// +/// - a binary operator that returns NULL on NULL input, applied directly to the +/// column (e.g. `a = 10`, `a < b`); +/// - an `IS NOT NULL` check on the column (e.g. `a IS NOT NULL`). +/// +/// This analysis is conservative; for example, OR clauses are not considered +/// null-rejecting, and neither are indirect operands like `a + 1 < 10`. +fn collect_null_rejecting_columns(predicate: &Arc<dyn PhysicalExpr>) -> HashSet<usize> { + let mut columns = HashSet::new(); + + for expr in split_conjunction(predicate) { + // `col IS NOT NULL` keeps only rows where `col` is non-null. + if let Some(is_not_null) = expr.downcast_ref::<IsNotNullExpr>() { + if let Some(col) = is_not_null.arg().downcast_ref::<Column>() { + columns.insert(col.index()); + } + continue; + } + + // A binary operator that returns NULL on NULL input rejects rows where + // a direct column operand is NULL. + if let Some(binary) = expr.downcast_ref::<BinaryExpr>() { + if !binary.op().returns_null_on_null() { + continue; + } + if let Some(col) = binary.left().downcast_ref::<Column>() { + columns.insert(col.index()); + } + if let Some(col) = binary.right().downcast_ref::<Column>() { + columns.insert(col.index()); + } + } + } + + columns +} + +/// Converts an interval bound to a [`Precision`] value. NULL bounds (which +/// represent "unbounded" in the interval type) map to [`Precision::Absent`]. +fn interval_bound_to_precision( + bound: ScalarValue, + is_exact: bool, +) -> Precision<ScalarValue> { + if bound.is_null() { + Precision::Absent + } else if is_exact { + Precision::Exact(bound) + } else { + Precision::Inexact(bound) + } +} + +/// Caps a row-bounded column statistic (a null count or distinct count) at the +/// filtered row count, since a column cannot have more nulls or distinct values +/// than it has rows. Known counts are demoted to inexact because a +/// filter-derived row bound is normally an estimate, the exception being an +/// exact zero, which proves the column is empty. +fn cap_at_rows( + value: Precision<usize>, + filtered_num_rows: Precision<usize>, +) -> Precision<usize> { + match filtered_num_rows { + Precision::Absent => value.to_inexact(), + Precision::Exact(0) => Precision::Exact(0), + rows => value.to_inexact().min(&rows), + } +} + +/// Scales a byte size by the filter selectivity. An exact zero row count means +/// the output is exactly empty, so the byte size is an exact zero too. +fn scale_byte_size_at_rows( + byte_size: Precision<usize>, + selectivity: f64, + filtered_num_rows: Precision<usize>, +) -> Precision<usize> { + if filtered_num_rows == Precision::Exact(0) { + Precision::Exact(0) + } else { + byte_size.with_estimated_selectivity(selectivity) + } +} + +/// Returns the NDV for a column constrained to one non-null value (e.g. +/// `column = literal` or a singleton interval), derived from the filtered row +/// estimate: zero rows means zero distinct values, a known positive row count +/// means exactly one, and an unknown row count means an inexact one (the column +/// could still be empty). +/// +/// The caller is responsible for proving the singleton domain. +fn distinct_count_for_singleton_domain( + filtered_num_rows: Precision<usize>, +) -> Precision<usize> { + match filtered_num_rows { + Precision::Exact(0) | Precision::Inexact(0) => filtered_num_rows, + // The row count is unknown, so the column could still be empty (zero + // distinct values); report an inexact one rather than overstating it. + Precision::Absent => Precision::Inexact(1), + _ => Precision::Exact(1), + } +} + +/// Builds output column statistics from interval-analysis boundaries. +/// +/// The interval bounds become min/max values, singleton intervals become +/// singleton NDV, and row-bounded counts are kept consistent with the filtered +/// row estimate. +fn collect_new_statistics( + schema: &SchemaRef, + input_column_stats: &[ColumnStatistics], + analysis_boundaries: Vec<ExprBoundaries>, + selectivity: f64, + null_rejecting_columns: &HashSet<usize>, + filtered_num_rows: Precision<usize>, +) -> Vec<ColumnStatistics> { + analysis_boundaries + .into_iter() + .enumerate() + .map( + |( + idx, + ExprBoundaries { + interval, + distinct_count, + .. + }, + )| { + let Some(interval) = interval else { + // If the interval is `None`, we can say that there are no rows. + // Use a typed null to preserve the column's data type, so that + // downstream interval analysis can still intersect intervals + // of the same type. + let typed_null = ScalarValue::try_from(schema.field(idx).data_type()) + .unwrap_or(ScalarValue::Null); + return ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Exact(typed_null.clone()), + min_value: Precision::Exact(typed_null.clone()), + sum_value: Precision::Exact(typed_null), + distinct_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + }; + }; + let (lower, upper) = interval.into_bounds(); + let is_single_value = + !lower.is_null() && !upper.is_null() && lower == upper; + let min_value = interval_bound_to_precision(lower, is_single_value); + let max_value = interval_bound_to_precision(upper, is_single_value); + + // Distinct and null counts cannot exceed the number of rows + // that survive the filter. Singleton intervals and + // null-rejecting predicates provide tighter bounds. + let capped_distinct_count = if is_single_value { + distinct_count_for_singleton_domain(filtered_num_rows) + } else { + cap_at_rows(distinct_count, filtered_num_rows) + }; + let capped_null_count = if null_rejecting_columns.contains(&idx) { + Precision::Exact(0) + } else { + cap_at_rows(input_column_stats[idx].null_count, filtered_num_rows) + }; + let byte_size = scale_byte_size_at_rows( + input_column_stats[idx].byte_size, + selectivity, + filtered_num_rows, + ); + ColumnStatistics { + null_count: capped_null_count, + max_value, + min_value, + sum_value: Precision::Absent, + distinct_count: capped_distinct_count, + byte_size, + } + }, + ) + .collect() +} + +/// The FilterExec streams wraps the input iterator and applies the predicate expression to +/// determine which rows to include in its output batches +struct FilterExecStream { + /// Output schema after the projection + schema: SchemaRef, + /// The expression to filter on. This expression must evaluate to a boolean value. + predicate: Arc<dyn PhysicalExpr>, + /// The input partition to filter. + input: SendableRecordBatchStream, + /// Runtime metrics recording + metrics: FilterExecMetrics, + /// The projection indices of the columns in the input schema + projection: Option<ProjectionRef>, + /// Batch coalescer to combine small batches + batch_coalescer: LimitedBatchCoalescer, +} + +/// The metrics for `FilterExec` +struct FilterExecMetrics { + /// Common metrics for most operators + baseline_metrics: BaselineMetrics, + /// Selectivity of the filter, calculated as output_rows / input_rows + selectivity: RatioMetrics, + // Remember to update `docs/source/user-guide/metrics.md` when adding new metrics, + // or modifying metrics comments +} + +impl FilterExecMetrics { + pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + baseline_metrics: BaselineMetrics::new(metrics, partition), + selectivity: MetricBuilder::new(metrics) + .with_type(MetricType::Summary) + .ratio_metrics("selectivity", partition), + } + } +} + +pub fn batch_filter( + batch: &RecordBatch, + predicate: &Arc<dyn PhysicalExpr>, +) -> Result<RecordBatch> { + filter_and_project(batch, predicate, None) +} + +fn filter_and_project( + batch: &RecordBatch, + predicate: &Arc<dyn PhysicalExpr>, + projection: Option<&Vec<usize>>, +) -> Result<RecordBatch> { + predicate + .evaluate(batch) + .and_then(|v| v.into_array(batch.num_rows())) + .and_then(|array| { + Ok(match (as_boolean_array(&array), projection) { + // Apply filter array to record batch + (Ok(filter_array), None) => filter_record_batch(batch, filter_array)?, + (Ok(filter_array), Some(projection)) => { + let projected_batch = batch.project(projection)?; + filter_record_batch(&projected_batch, filter_array)? + } + (Err(_), _) => { + return internal_err!( + "Cannot create filter_array from non-boolean predicates" + ); + } + }) + }) +} + +impl Stream for FilterExecStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let elapsed_compute = self.metrics.baseline_metrics.elapsed_compute().clone(); + loop { + // If there is a completed batch ready, return it + if let Some(batch) = self.batch_coalescer.next_completed_batch() { + self.metrics.selectivity.add_part(batch.num_rows()); + let poll = Poll::Ready(Some(Ok(batch))); + return self.metrics.baseline_metrics.record_poll(poll); + } + + if self.batch_coalescer.is_finished() { + // If input is done and no batches are ready, return None to signal end of stream. + return Poll::Ready(None); + } + + // Attempt to pull the next batch from the input stream. + match ready!(self.input.poll_next_unpin(cx)) { + None => { + self.batch_coalescer.finish()?; + // Release the input pipeline's resources. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + // continue draining the coalescer + } + Some(Ok(batch)) => { + let timer = elapsed_compute.timer(); + let status = self.predicate.as_ref() + .evaluate(&batch) + .and_then(|v| v.into_array(batch.num_rows())) + .and_then(|array| { + Ok(match self.projection.as_ref() { + Some(projection) => { + let projected_batch = batch.project(projection)?; + (array, projected_batch) + }, + None => (array, batch) + }) + }).and_then(|(array, batch)| { + match as_boolean_array(&array) { + Ok(filter_array) => { + self.metrics.selectivity.add_total(batch.num_rows()); + // TODO: support push_batch_with_filter in LimitedBatchCoalescer + let batch = filter_record_batch(&batch, filter_array)?; + let state = self.batch_coalescer.push_batch(batch)?; + Ok(state) + } + Err(_) => { + internal_err!( + "Cannot create filter_array from non-boolean predicates" + ) + } + } + })?; + timer.done(); + + match status { + PushBatchStatus::Continue => { + // Keep pushing more batches + } + PushBatchStatus::LimitReached => { + // limit was reached, so stop early + self.batch_coalescer.finish()?; + // Release the input pipeline's resources. + let input_schema = self.input.schema(); + self.input = + Box::pin(EmptyRecordBatchStream::new(input_schema)); + // continue draining the coalescer + } + } + } + + // Error case + other => return Poll::Ready(other), + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + // Same number of record batches + self.input.size_hint() + } +} +impl RecordBatchStream for FilterExecStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Return the equals Column-Pairs and Non-equals Column-Pairs +#[deprecated( + since = "51.0.0", + note = "This function will be internal in the future" +)] +pub fn collect_columns_from_predicate( + predicate: &'_ Arc<dyn PhysicalExpr>, +) -> EqualAndNonEqual<'_> { + collect_columns_from_predicate_inner(predicate) +} + +fn collect_columns_from_predicate_inner( + predicate: &'_ Arc<dyn PhysicalExpr>, +) -> EqualAndNonEqual<'_> { + let mut eq_predicate_columns = Vec::<PhysicalExprPairRef>::new(); + let mut ne_predicate_columns = Vec::<PhysicalExprPairRef>::new(); + + let predicates = split_conjunction(predicate); + predicates.into_iter().for_each(|p| { + if let Some(binary) = p.downcast_ref::<BinaryExpr>() { + // Only extract pairs where at least one side is a Column reference. + // Pairs like `complex_expr = literal` should not create equivalence + // classes — the literal could appear in many unrelated expressions + // (e.g. sort keys), and normalize_expr's deep traversal would + // replace those occurrences with the complex expression, corrupting + // sort orderings. Constant propagation for such pairs is handled + // separately by `extend_constants`. + let has_direct_column_operand = + binary.left().downcast_ref::<Column>().is_some() + || binary.right().downcast_ref::<Column>().is_some(); + if !has_direct_column_operand { + return; + } + match binary.op() { + Operator::Eq => { + eq_predicate_columns.push((binary.left(), binary.right())) + } + Operator::NotEq => { + ne_predicate_columns.push((binary.left(), binary.right())) + } + _ => {} + } + } + }); + + (eq_predicate_columns, ne_predicate_columns) +} + +/// Pair of `Arc<dyn PhysicalExpr>`s +pub type PhysicalExprPairRef<'a> = (&'a Arc<dyn PhysicalExpr>, &'a Arc<dyn PhysicalExpr>); + +/// The equals Column-Pairs and Non-equals Column-Pairs in the Predicates +pub type EqualAndNonEqual<'a> = + (Vec<PhysicalExprPairRef<'a>>, Vec<PhysicalExprPairRef<'a>>); + +#[cfg(test)] +mod tests { + use super::*; + use crate::empty::EmptyExec; + use crate::expressions::*; + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::test; + use crate::test::exec::StatisticsExec; + use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; + + #[tokio::test] + async fn collect_columns_predicates() -> Result<()> { + let schema = test::aggr_test_schema(); + let predicate: Arc<dyn PhysicalExpr> = binary( + binary( + binary(col("c2", &schema)?, Operator::GtEq, lit(1u32), &schema)?, + Operator::And, + binary(col("c2", &schema)?, Operator::Eq, lit(4u32), &schema)?, + &schema, + )?, + Operator::And, + binary( + binary( + col("c2", &schema)?, + Operator::Eq, + col("c9", &schema)?, + &schema, + )?, + Operator::And, + binary( + col("c1", &schema)?, + Operator::NotEq, + col("c13", &schema)?, + &schema, + )?, + &schema, + )?, + &schema, + )?; + + let (equal_pairs, ne_pairs) = collect_columns_from_predicate_inner(&predicate); + assert_eq!(2, equal_pairs.len()); + assert!(equal_pairs[0].0.eq(&col("c2", &schema)?)); + assert!(equal_pairs[0].1.eq(&lit(4u32))); + + assert!(equal_pairs[1].0.eq(&col("c2", &schema)?)); + assert!(equal_pairs[1].1.eq(&col("c9", &schema)?)); + + assert_eq!(1, ne_pairs.len()); + assert!(ne_pairs[0].0.eq(&col("c1", &schema)?)); + assert!(ne_pairs[0].1.eq(&col("c13", &schema)?)); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_basic_expr() -> Result<()> { + // Table: + // a: min=1, max=100 + let bytes_per_row = 4; + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(100 * bytes_per_row), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }], + }, + schema.clone(), + )); + + // a <= 25 + let predicate: Arc<dyn PhysicalExpr> = + binary(col("a", &schema)?, Operator::LtEq, lit(25i32), &schema)?; + + // WHERE a <= 25 + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(25)); + assert_eq!( + statistics.total_byte_size, + Precision::Inexact(25 * bytes_per_row) + ); + assert_eq!( + statistics.column_statistics, + vec![ColumnStatistics { + // `a <= 25` rejects nulls, so the column has no surviving nulls. + null_count: Precision::Exact(0), + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(25))), + ..Default::default() + }] + ); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_column_level_nested() -> Result<()> { + // Table: + // a: min=1, max=100 + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }], + total_byte_size: Precision::Absent, + }, + schema.clone(), + )); + + // WHERE a <= 25 + let sub_filter: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new( + binary(col("a", &schema)?, Operator::LtEq, lit(25i32), &schema)?, + input, + )?); + + // Nested filters (two separate physical plans, instead of AND chain in the expr) + // WHERE a >= 10 + // WHERE a <= 25 + let filter: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new( + binary(col("a", &schema)?, Operator::GtEq, lit(10i32), &schema)?, + sub_filter, + )?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(16)); + assert_eq!( + statistics.column_statistics, + vec![ColumnStatistics { + // `a <= 25 AND a >= 10` rejects nulls in `a`. + null_count: Precision::Exact(0), + min_value: Precision::Inexact(ScalarValue::Int32(Some(10))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(25))), + ..Default::default() + }] + ); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_column_level_nested_multiple() -> Result<()> { + // Table: + // a: min=1, max=100 + // b: min=1, max=50 + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(50))), + ..Default::default() + }, + ], + total_byte_size: Precision::Absent, + }, + schema.clone(), + )); + + // WHERE a <= 25 + let a_lte_25: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new( + binary(col("a", &schema)?, Operator::LtEq, lit(25i32), &schema)?, + input, + )?); + + // WHERE b > 45 + let b_gt_5: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new( + binary(col("b", &schema)?, Operator::Gt, lit(45i32), &schema)?, + a_lte_25, + )?); + + // WHERE a >= 10 + let filter: Arc<dyn ExecutionPlan> = Arc::new(FilterExec::try_new( + binary(col("a", &schema)?, Operator::GtEq, lit(10i32), &schema)?, + b_gt_5, + )?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + // On a uniform distribution, only fifteen rows will satisfy the + // filter that 'a' proposed (a >= 10 AND a <= 25) (15/100) and only + // 5 rows will satisfy the filter that 'b' proposed (b > 45) (5/50). + // + // Which would result with a selectivity of '15/100 * 5/50' or 0.015 + // and that means about %1.5 of the all rows (rounded up to 2 rows). + assert_eq!(statistics.num_rows, Precision::Inexact(2)); + assert_eq!( + statistics.column_statistics, + vec![ + ColumnStatistics { + // `a <= 25 AND a >= 10` rejects nulls in `a`. + null_count: Precision::Exact(0), + min_value: Precision::Inexact(ScalarValue::Int32(Some(10))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(25))), + ..Default::default() + }, + ColumnStatistics { + // `b > 45` in the upstream filter zeroes b's nulls; the outer + // filter then caps the (already zero) count, demoting to inexact. + null_count: Precision::Inexact(0), + min_value: Precision::Inexact(ScalarValue::Int32(Some(46))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(50))), + ..Default::default() + } + ] + ); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_when_input_stats_missing() -> Result<()> { + // Table: + // a: min=???, max=??? (missing) + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema.clone(), + )); + + // a <= 25 + let predicate: Arc<dyn PhysicalExpr> = + binary(col("a", &schema)?, Operator::LtEq, lit(25i32), &schema)?; + + // WHERE a <= 25 + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Absent); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_multiple_columns() -> Result<()> { + // Table: + // a: min=1, max=100 + // b: min=1, max=3 + // c: min=1000.0 max=1100.0 + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Float32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(3))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Float32(Some(1000.0))), + max_value: Precision::Inexact(ScalarValue::Float32(Some(1100.0))), + ..Default::default() + }, + ], + }, + schema, + )); + // WHERE a<=53 AND (b=3 AND (c<=1075.0 AND a>b)) + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::LtEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(53)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(3)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::LtEq, + Arc::new(Literal::new(ScalarValue::Float32(Some(1075.0)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Column::new("b", 1)), + )), + )), + )), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + // 0.5 (from a) * 0.333333... (from b) * 0.798387... (from c) ≈ 0.1330... + // num_rows after ceil => 133.0... => 134 + // total_byte_size after ceil => 532.0... => 533 + assert_eq!(statistics.num_rows, Precision::Inexact(134)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(533)); + let exp_col_stats = vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(4))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(53))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(3))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(3))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Float32(Some(1000.0))), + max_value: Precision::Inexact(ScalarValue::Float32(Some(1075.0))), + ..Default::default() + }, + ]; + let _ = exp_col_stats + .into_iter() + .zip(statistics.column_statistics.clone()) + .map(|(expected, actual)| { + if let Some(val) = actual.min_value.get_value() { + if val.data_type().is_floating() { + // Windows rounds arithmetic operation results differently for floating point numbers. + // Therefore, we check if the actual values are in an epsilon range. + let actual_min = actual.min_value.get_value().unwrap(); + let actual_max = actual.max_value.get_value().unwrap(); + let expected_min = expected.min_value.get_value().unwrap(); + let expected_max = expected.max_value.get_value().unwrap(); + let eps = ScalarValue::Float32(Some(1e-6)); + + assert!(actual_min.sub(expected_min).unwrap() < eps); + assert!(actual_min.sub(expected_min).unwrap() < eps); + + assert!(actual_max.sub(expected_max).unwrap() < eps); + assert!(actual_max.sub(expected_max).unwrap() < eps); + } else { + assert_eq!(actual, expected); + } + } else { + assert_eq!(actual, expected); + } + }); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_full_selective() -> Result<()> { + // Table: + // a: min=1, max=100 + // b: min=1, max=3 + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(3))), + ..Default::default() + }, + ], + }, + schema, + )); + // WHERE a<200 AND 1<=b + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(200)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + Operator::LtEq, + Arc::new(Column::new("b", 1)), + )), + )); + // The filter predicate passes all (non-null) entries, so min/max/NDV + // are unchanged. `a < 200` and `1 <= b` are null-rejecting, though, so + // both columns lose any nulls regardless of selectivity. + let mut expected = StatisticsContext::new() + .compute(input.as_ref(), &StatisticsArgs::new())? + .column_statistics + .clone(); + for col in &mut expected { + col.null_count = Precision::Exact(0); + } + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Inexact(1000)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(4000)); + assert_eq!(statistics.column_statistics, expected); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_zero_selective() -> Result<()> { + // Table: + // a: min=1, max=100 + // b: min=1, max=3 + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(3))), + ..Default::default() + }, + ], + }, + schema, + )); + // WHERE a>200 AND 1<=b + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(200)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + Operator::LtEq, + Arc::new(Column::new("b", 1)), + )), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Inexact(0)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(0)); + assert_eq!( + statistics.column_statistics, + vec![ + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(None)), + max_value: Precision::Exact(ScalarValue::Int32(None)), + sum_value: Precision::Exact(ScalarValue::Int32(None)), + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + }, + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(None)), + max_value: Precision::Exact(ScalarValue::Int32(None)), + sum_value: Precision::Exact(ScalarValue::Int32(None)), + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + }, + ] + ); + + Ok(()) + } + + /// Regression test: stacking two FilterExecs where the inner filter + /// proves zero selectivity should not panic with a type mismatch + /// during interval intersection. + /// + /// Previously, when a filter proved no rows could match, the column + /// statistics used untyped `ScalarValue::Null` (data type `Null`). + /// If an outer FilterExec then tried to analyze its own predicate + /// against those statistics, `Interval::intersect` would fail with: + /// "Only intervals with the same data type are intersectable, lhs:Null, rhs:Int32" + #[tokio::test] + async fn test_nested_filter_with_zero_selectivity_inner() -> Result<()> { + // Inner table: a: [1, 100], b: [1, 3] + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(3))), + ..Default::default() + }, + ], + }, + schema, + )); + + // Inner filter: a > 200 (impossible given a max=100 → zero selectivity) + let inner_predicate: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(200)))), + )); + let inner_filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(inner_predicate, input)?); + + // Outer filter: a = 50 + // Before the fix, this would panic because the inner filter's + // zero-selectivity statistics produced Null-typed intervals for + // column `a`, which couldn't intersect with the Int32 literal. + let outer_predicate: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(50)))), + )); + let outer_filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(outer_predicate, inner_filter)?); + + // Should succeed without error + let statistics = StatisticsContext::new() + .compute(outer_filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(0)); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_more_inputs() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ], + }, + schema, + )); + // WHERE a<50 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(50)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Inexact(490)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(1960)); + assert_eq!( + statistics.column_statistics, + vec![ + ColumnStatistics { + // `a < 50` rejects nulls in `a`. + null_count: Precision::Exact(0), + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(49))), + ..Default::default() + }, + // `b` is not referenced by the predicate, so its stats are + // unchanged (null count stays unknown). + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ] + ); + + Ok(()) + } + + #[tokio::test] + async fn test_empty_input_statistics() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema, + )); + // WHERE a <= 10 AND 0 <= a - 5 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::LtEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + Operator::LtEq, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Minus, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )), + )), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let filter_statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + let expected_filter_statistics = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![ColumnStatistics { + // `a <= 10` rejects nulls, so `a` has no surviving nulls even + // though the input statistics are entirely unknown. + null_count: Precision::Exact(0), + min_value: Precision::Inexact(ScalarValue::Int32(Some(5))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(10))), + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + byte_size: Precision::Absent, + }], + }; + + assert_eq!(*filter_statistics, expected_filter_statistics); + + Ok(()) + } + + #[tokio::test] + async fn test_statistics_with_constant_column() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema, + )); + // WHERE a = 10 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let filter_statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + // First column is "a", and it is a column with only one value after the filter. + assert!(filter_statistics.column_statistics[0].is_singleton()); + + Ok(()) + } + + #[tokio::test] + async fn test_validation_filter_selectivity() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema, + )); + // WHERE a = 10 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )); + let filter = FilterExec::try_new(predicate, input)?; + assert!(filter.with_default_selectivity(120).is_err()); + Ok(()) + } + + #[tokio::test] + async fn test_custom_filter_selectivity() -> Result<()> { + // Need a decimal to trigger inexact selectivity + let schema = + Schema::new(vec![Field::new("a", DataType::Decimal128(2, 3), false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ColumnStatistics { + ..Default::default() + }], + }, + schema, + )); + // WHERE a = 10 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Decimal128(Some(10), 10, 10))), + )); + let filter = FilterExec::try_new(predicate, input)?; + let statistics = + StatisticsContext::new().compute(&filter, &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(200)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(800)); + let filter = filter.with_default_selectivity(40)?; + let statistics = + StatisticsContext::new().compute(&filter, &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(400)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(1600)); + Ok(()) + } + + #[test] + fn test_equivalence_properties_union_type() -> Result<()> { + let union_type = DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("f1", DataType::Int32, true), + Field::new("f2", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Sparse, + ); + + let schema = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, true), + Field::new("c2", union_type, true), + ])); + + let exec = FilterExec::try_new( + binary( + binary(col("c1", &schema)?, Operator::GtEq, lit(1i32), &schema)?, + Operator::And, + binary(col("c1", &schema)?, Operator::LtEq, lit(4i32), &schema)?, + &schema, + )?, + Arc::new(EmptyExec::new(Arc::clone(&schema))), + )?; + + StatisticsContext::new() + .compute(&exec, &StatisticsArgs::new()) + .unwrap(); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_with_projection() -> Result<()> { + // Create a schema with multiple columns + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ])); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + // Create a filter predicate: a > 10 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )); + + // Create filter with projection [0, 2] (columns a and c) using builder + let projection = Some(vec![0, 2]); + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(projection.clone()) + .unwrap() + .build()?; + + // Verify projection is set correctly + assert_eq!(filter.projection(), &Some([0, 2].into())); + + // Verify schema contains only projected columns + let output_schema = filter.schema(); + assert_eq!(output_schema.fields().len(), 2); + assert_eq!(output_schema.field(0).name(), "a"); + assert_eq!(output_schema.field(1).name(), "c"); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_without_projection() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + // Create filter without projection using builder + let filter = FilterExecBuilder::new(predicate, input).build()?; + + // Verify no projection is set + assert!(filter.projection().is_none()); + + // Verify schema contains all columns + let output_schema = filter.schema(); + assert_eq!(output_schema.fields().len(), 2); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_invalid_projection() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + // Try to create filter with invalid projection (index out of bounds) using builder + let result = + FilterExecBuilder::new(predicate, input).apply_projection(Some(vec![0, 5])); // 5 is out of bounds + + // Should return an error + assert!(result.is_err()); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_vs_with_projection() -> Result<()> { + // This test verifies that the builder with projection produces the same result + // as try_new().with_projection(), but more efficiently (one compute_properties call) + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + ]); + + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(4000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + ..Default::default() + }, + ColumnStatistics { + ..Default::default() + }, + ColumnStatistics { + ..Default::default() + }, + ], + }, + schema, + )); + let input: Arc<dyn ExecutionPlan> = input; + + let predicate: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(50)))), + )); + + let projection = Some(vec![0, 2]); + + // Method 1: Builder with projection (one call to compute_properties) + let filter1 = FilterExecBuilder::new(Arc::clone(&predicate), Arc::clone(&input)) + .apply_projection(projection.clone()) + .unwrap() + .build()?; + + // Method 2: Also using builder for comparison (deprecated try_new().with_projection() removed) + let filter2 = FilterExecBuilder::new(predicate, input) + .apply_projection(projection) + .unwrap() + .build()?; + + // Both methods should produce equivalent results + assert_eq!(filter1.schema(), filter2.schema()); + assert_eq!(filter1.projection(), filter2.projection()); + + // Verify statistics are the same + let stats1 = + StatisticsContext::new().compute(&filter1, &StatisticsArgs::new())?; + let stats2 = + StatisticsContext::new().compute(&filter2, &StatisticsArgs::new())?; + assert_eq!(stats1.num_rows, stats2.num_rows); + assert_eq!(stats1.total_byte_size, stats2.total_byte_size); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_statistics_with_projection() -> Result<()> { + // Test that statistics are correctly computed when using builder with projection + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(12000), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(10))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(200))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(5))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(50))), + ..Default::default() + }, + ], + }, + schema, + )); + + // Filter: a < 50, Project: [0, 2] + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(50)))), + )); + + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![0, 2])) + .unwrap() + .build()?; + + let statistics = + StatisticsContext::new().compute(&filter, &StatisticsArgs::new())?; + + // Verify statistics reflect both filtering and projection + assert!(matches!(statistics.num_rows, Precision::Inexact(_))); + + // Schema should only have 2 columns after projection + assert_eq!(filter.schema().fields().len(), 2); + + Ok(()) + } + + #[test] + fn test_builder_predicate_validation() -> Result<()> { + // Test that builder validates predicate type correctly + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + // Create a predicate that doesn't return boolean (returns Int32) + let invalid_predicate = Arc::new(Column::new("a", 0)); + + // Should fail because predicate doesn't return boolean + let result = FilterExecBuilder::new(invalid_predicate, input) + .apply_projection(Some(vec![0])) + .unwrap() + .build(); + + assert!(result.is_err()); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_projection_composition() -> Result<()> { + // Test that calling apply_projection multiple times composes projections + // If initial projection is [0, 2, 3] and we call apply_projection([0, 2]), + // the result should be [0, 3] (indices 0 and 2 of [0, 2, 3]) + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + ])); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + // Create a filter predicate: a > 10 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )); + + // First projection: [0, 2, 3] -> select columns a, c, d + // Second projection: [0, 2] -> select indices 0 and 2 of [0, 2, 3] -> [0, 3] + // Final result: columns a and d + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![0, 2, 3]))? + .apply_projection(Some(vec![0, 2]))? + .build()?; + + // Verify composed projection is [0, 3] + assert_eq!(filter.projection(), &Some([0, 3].into())); + + // Verify schema contains only columns a and d + let output_schema = filter.schema(); + assert_eq!(output_schema.fields().len(), 2); + assert_eq!(output_schema.field(0).name(), "a"); + assert_eq!(output_schema.field(1).name(), "d"); + + Ok(()) + } + + #[tokio::test] + async fn test_builder_projection_composition_none_clears() -> Result<()> { + // Test that passing None clears the projection + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )); + + // Set a projection then clear it with None + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![0]))? + .apply_projection(None)? + .build()?; + + // Projection should be cleared + assert_eq!(filter.projection(), &None); + + // Schema should have all columns + let output_schema = filter.schema(); + assert_eq!(output_schema.fields().len(), 2); + + Ok(()) + } + + #[test] + fn test_filter_with_projection_remaps_post_phase_parent_filters() -> Result<()> { + // Test that FilterExec with a projection must remap parent dynamic + // filter column indices from its output schema to the input schema + // before passing them to the child. + let input_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::Float64, false), + ])); + let input = Arc::new(EmptyExec::new(Arc::clone(&input_schema))); + + // FilterExec: a > 0, projection=[c@2] + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + )); + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![2]))? + .build()?; + + // Output schema should be [c:Float64] + let output_schema = filter.schema(); + assert_eq!(output_schema.fields().len(), 1); + assert_eq!(output_schema.field(0).name(), "c"); + + // Simulate a parent dynamic filter referencing output column c@0 + let parent_filter: Arc<dyn PhysicalExpr> = Arc::new(Column::new("c", 0)); + + let config = ConfigOptions::new(); + let desc = filter.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![parent_filter], + &config, + )?; + + // The filter pushed to the child must reference c@2 (input schema), + // not c@0 (output schema). + let parent_filters = desc.parent_filters(); + assert_eq!(parent_filters.len(), 1); // one child + assert_eq!(parent_filters[0].len(), 1); // one filter + let remapped = &parent_filters[0][0].predicate; + let display = format!("{remapped}"); + assert_eq!( + display, "c@2", + "Post-phase parent filter column index must be remapped \ + from output schema (c@0) to input schema (c@2)" + ); + + Ok(()) + } + + /// Regression test for https://github.com/apache/datafusion/issues/20194 + /// + /// `collect_columns_from_predicate_inner` should only extract equality + /// pairs where at least one side is a Column. Pairs like + /// `complex_expr = literal` must not create equivalence classes because + /// `normalize_expr`'s deep traversal would replace the literal inside + /// unrelated expressions (e.g. sort keys) with the complex expression. + #[test] + fn test_collect_columns_skips_non_column_pairs() -> Result<()> { + let schema = test::aggr_test_schema(); + + // Simulate: nvl(c2, 0) = 0 → (c2 IS DISTINCT FROM 0) = 0 + // Neither side is a Column, so this should NOT be extracted. + let complex_expr: Arc<dyn PhysicalExpr> = binary( + col("c2", &schema)?, + Operator::IsDistinctFrom, + lit(0u32), + &schema, + )?; + let predicate: Arc<dyn PhysicalExpr> = + binary(complex_expr, Operator::Eq, lit(0u32), &schema)?; + + let (equal_pairs, _) = collect_columns_from_predicate_inner(&predicate); + assert_eq!( + 0, + equal_pairs.len(), + "Should not extract equality pairs where neither side is a Column" + ); + + // But col = literal should still be extracted + let predicate: Arc<dyn PhysicalExpr> = + binary(col("c2", &schema)?, Operator::Eq, lit(0u32), &schema)?; + let (equal_pairs, _) = collect_columns_from_predicate_inner(&predicate); + assert_eq!( + 1, + equal_pairs.len(), + "Should extract equality pairs where one side is a Column" + ); + + Ok(()) + } + + /// Columns with Absent min/max statistics should remain Absent after + /// FilterExec. + #[tokio::test] + async fn test_filter_statistics_absent_columns_stay_absent() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Absent, + column_statistics: vec![ + ColumnStatistics::default(), + ColumnStatistics::default(), + ], + }, + schema.clone(), + )); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + let col_b_stats = &statistics.column_statistics[1]; + assert_eq!(col_b_stats.min_value, Precision::Absent); + assert_eq!(col_b_stats.max_value, Precision::Absent); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_ndv() -> Result<()> { + #[expect(clippy::type_complexity)] + let cases: Vec<( + &str, + Vec<Field>, + Vec<ColumnStatistics>, + Arc<dyn PhysicalExpr>, + Vec<Precision<usize>>, + )> = vec![ + ( + "utf8 equality", + vec![Field::new("name", DataType::Utf8, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("hello".to_string())))), + )), + vec![Precision::Exact(1)], + ), + ( + "utf8view equality", + vec![Field::new("name", DataType::Utf8View, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8View(Some( + "hello".to_string(), + )))), + )), + vec![Precision::Exact(1)], + ), + ( + "largeutf8 equality", + vec![Field::new("name", DataType::LargeUtf8, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::LargeUtf8(Some( + "hello".to_string(), + )))), + )), + vec![Precision::Exact(1)], + ), + ( + "utf8 reversed (literal = column)", + vec![Field::new("name", DataType::Utf8, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Utf8(Some("hello".to_string())))), + Operator::Eq, + Arc::new(Column::new("name", 0)), + )), + vec![Precision::Exact(1)], + ), + ( + "OR is not collapsed to NDV=1, but NDV is capped at filtered rows", + vec![Field::new("name", DataType::Utf8, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("a".to_string())))), + )), + Operator::Or, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("b".to_string())))), + )), + )), + // Input NDV is 50, but the 20% default selectivity on 100 rows + // estimates 20 output rows, so NDV is capped at 20. + vec![Precision::Inexact(20)], + ), + ( + "AND with mixed types (Utf8 + Int32)", + vec![ + Field::new("name", DataType::Utf8, false), + Field::new("age", DataType::Int32, false), + ], + vec![ + ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }, + ColumnStatistics { + distinct_count: Precision::Inexact(80), + ..Default::default() + }, + ], + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "hello".to_string(), + )))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("age", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + )), + vec![Precision::Exact(1), Precision::Exact(1)], + ), + ( + "numeric equality with min/max bounds (interval analysis path)", + vec![Field::new("a", DataType::Int32, false)], + vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + distinct_count: Precision::Inexact(80), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + vec![Precision::Exact(1)], + ), + ( + "timestamp equality", + vec![Field::new( + "ts", + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + false, + )], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(500), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Column::new("ts", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::TimestampNanosecond( + Some(1_609_459_200_000_000_000), + None, + ))), + )), + vec![Precision::Exact(1)], + ), + ( + "contradictory numeric equality (infeasible)", + vec![Field::new("a", DataType::Int32, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(99)))), + )), + )), + vec![Precision::Exact(0)], + ), + ( + "utf8 equality with absent input NDV", + vec![Field::new("name", DataType::Utf8, false)], + vec![ColumnStatistics { + distinct_count: Precision::Absent, + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("hello".to_string())))), + )), + vec![Precision::Exact(1)], + ), + ( + "contradictory utf8 equality (infeasible)", + vec![Field::new("name", DataType::Utf8, false)], + vec![ColumnStatistics { + distinct_count: Precision::Inexact(100), + ..Default::default() + }], + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "alice".to_string(), + )))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "bob".to_string(), + )))), + )), + )), + vec![Precision::Exact(0)], + ), + ( + "redundant same-value equality combined with another column", + vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ], + vec![ + ColumnStatistics { + distinct_count: Precision::Inexact(80), + ..Default::default() + }, + ColumnStatistics { + distinct_count: Precision::Inexact(40), + ..Default::default() + }, + ], + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )), + )), + vec![Precision::Exact(1), Precision::Exact(1)], + ), + ]; + + for (desc, fields, col_stats, predicate, expected_ndvs) in cases { + let schema = Schema::new(fields); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: col_stats, + }, + schema.clone(), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = StatisticsContext::new() + .compute(filter.as_ref(), &StatisticsArgs::new())?; + + for (i, expected) in expected_ndvs.iter().enumerate() { + assert_eq!( + statistics.column_statistics[i].distinct_count, *expected, + "case '{desc}': column {i} NDV mismatch" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_preserves_exactly_empty_input() -> Result<()> { + // A satisfiable predicate over an exactly empty input: the filter cannot + // produce rows, so the whole estimate stays exact. Column `b` is not + // mentioned by the predicate, so its null and distinct counts go through + // the generic row cap. + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]); + let input_stats = Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ + ColumnStatistics { + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + ..Default::default() + }, + ColumnStatistics { + null_count: Precision::Exact(3), + distinct_count: Precision::Exact(7), + byte_size: Precision::Exact(0), + ..Default::default() + }, + ], + }; + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + let input = Arc::new(StatisticsExec::new(input_stats, schema.clone())); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Exact(0)); + assert_eq!(statistics.total_byte_size, Precision::Exact(0)); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[1].null_count, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[1].distinct_count, + Precision::Exact(0) + ); + + // A contradictory predicate (`a = 1 AND a = 2`) discards all rows, the + // output is empty independently of the input. + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(8000), + column_statistics: vec![ColumnStatistics::new_unknown(); 2], + }, + schema, + )); + let contradiction = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(contradiction, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Exact(0)); + assert_eq!(statistics.total_byte_size, Precision::Exact(0)); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_exact_empty_input_zeroes_byte_size() -> Result<()> { + let cases = [ + ("absent", Precision::Absent, Precision::Absent), + ("inexact", Precision::Inexact(8000), Precision::Inexact(400)), + ]; + + for (desc, input_total_byte_size, input_byte_size) in cases { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let input_stats = Statistics { + num_rows: Precision::Exact(0), + total_byte_size: input_total_byte_size, + column_statistics: vec![ColumnStatistics { + byte_size: input_byte_size, + ..Default::default() + }], + }; + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + let input = Arc::new(StatisticsExec::new(input_stats, schema)); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = StatisticsContext::new() + .compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!( + statistics.num_rows, + Precision::Exact(0), + "case '{desc}': num_rows mismatch" + ); + assert_eq!( + statistics.total_byte_size, + Precision::Exact(0), + "case '{desc}': total_byte_size mismatch" + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Exact(0), + "case '{desc}': byte_size mismatch" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_empty_input_equality_ndv_zero() -> Result<()> { + let cases: Vec<(&str, Schema, Statistics, Arc<dyn PhysicalExpr>)> = vec![ + ( + "fallback string equality", + Schema::new(vec![Field::new("name", DataType::Utf8, true)]), + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + ..Default::default() + }], + }, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))), + )), + ), + ( + "interval numeric equality", + Schema::new(vec![Field::new("a", DataType::Int32, true)]), + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(10))), + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + ..Default::default() + }], + }, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )), + ), + ]; + + for (desc, schema, input_stats, predicate) in cases { + let input = Arc::new(StatisticsExec::new(input_stats, schema)); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = StatisticsContext::new() + .compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!( + statistics.num_rows, + Precision::Exact(0), + "case '{desc}': row count mismatch" + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(0), + "case '{desc}': NDV should be capped at zero rows" + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_and_equality_ndv() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1200), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(80), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(50))), + distinct_count: Precision::Inexact(40), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(200))), + null_count: Precision::Inexact(90), + distinct_count: Precision::Inexact(150), + ..Default::default() + }, + ], + }, + schema.clone(), + )); + + // a = 42 AND b > 10 AND c = 7 + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(7)))), + )), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + // Equality predicates collapse NDV and reject nulls for their columns. + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + // b > 10 narrows to [11, 50] but doesn't collapse to a single value. + // The combined selectivity of a=42 (1/80) and c=7 (1/150) on 100 rows + // computes num_rows = 1, so NDV is capped at the row count: min(40, 1) = 1. + assert_eq!( + statistics.column_statistics[1].distinct_count, + Precision::Inexact(1) + ); + assert_eq!( + statistics.column_statistics[2].distinct_count, + Precision::Exact(1) + ); + assert_eq!( + statistics.column_statistics[2].null_count, + Precision::Exact(0) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_absent_bounds_ndv() -> Result<()> { + // a: ndv=80, no min/max + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(400), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Inexact(80), + ..Default::default() + }], + }, + schema.clone(), + )); + + // Even without input bounds, interval analysis can derive singleton + // bounds from the equality itself. + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_int8_ndv() -> Result<()> { + // a: min=-100, max=100, ndv=50 + let schema = Schema::new(vec![Field::new("a", DataType::Int8, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(100), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int8(Some(-100))), + max_value: Precision::Inexact(ScalarValue::Int8(Some(100))), + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + }, + schema.clone(), + )); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int8(Some(42)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_int64_ndv() -> Result<()> { + // a: min=0, max=1_000_000, ndv=100_000 + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100_000), + total_byte_size: Precision::Inexact(800_000), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int64(Some(0))), + max_value: Precision::Inexact(ScalarValue::Int64(Some(1_000_000))), + distinct_count: Precision::Inexact(100_000), + ..Default::default() + }], + }, + schema.clone(), + )); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int64(Some(42)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_float32_ndv() -> Result<()> { + // a: min=0.0, max=100.0, ndv=50 + let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(400), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Float32(Some(0.0))), + max_value: Precision::Inexact(ScalarValue::Float32(Some(100.0))), + distinct_count: Precision::Inexact(50), + ..Default::default() + }], + }, + schema.clone(), + )); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Float32(Some(42.5)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_reversed_ndv() -> Result<()> { + // a: min=1, max=100, ndv=80 + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(400), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + distinct_count: Precision::Inexact(80), + ..Default::default() + }], + }, + schema.clone(), + )); + + // 42 = a (literal on the left) + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + Operator::Eq, + Arc::new(Column::new("a", 0)), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_equality_timestamp_ndv() -> Result<()> { + // ts: min=1_000_000_000, max=2_000_000_000, ndv=500 + let schema = Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + false, + )]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(8000), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::TimestampNanosecond( + Some(1_000_000_000), + None, + )), + max_value: Precision::Inexact(ScalarValue::TimestampNanosecond( + Some(2_000_000_000), + None, + )), + distinct_count: Precision::Inexact(500), + ..Default::default() + }], + }, + schema.clone(), + )); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("ts", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::TimestampNanosecond( + Some(1_500_000_000), + None, + ))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(1) + ); + Ok(()) + } + + #[test] + fn test_collect_equality_columns() { + use std::collections::HashSet; + // (description, predicate, expected_column_indices, expected_infeasible) + #[expect(clippy::type_complexity)] + let cases: Vec<(&str, Arc<dyn PhysicalExpr>, Vec<usize>, bool)> = vec![ + ( + "simple col = literal", + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + vec![0], + false, + ), + ( + "reversed literal = col", + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + Operator::Eq, + Arc::new(Column::new("a", 0)), + )), + vec![0], + false, + ), + ( + "AND with two equalities", + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "hello".to_string(), + )))), + )), + )), + vec![0, 1], + false, + ), + ( + "OR produces empty set", + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + Operator::Or, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(99)))), + )), + )), + vec![], + false, + ), + ( + "greater-than produces empty set", + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + vec![], + false, + ), + ( + "col = col produces empty set", + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Column::new("b", 1)), + )), + vec![], + false, + ), + ( + "nested AND with three equalities", + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 2)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(3)))), + )), + )), + vec![0, 1, 2], + false, + ), + ( + "AND with mixed equality and non-equality", + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )), + )), + vec![0], + false, + ), + ( + "col = NULL is excluded", + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(None))), + )), + vec![], + false, + ), + ( + "NULL = col is excluded", + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Utf8(None))), + Operator::Eq, + Arc::new(Column::new("a", 0)), + )), + vec![], + false, + ), + ( + "contradictory: same col, different literals", + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "alice".to_string(), + )))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "bob".to_string(), + )))), + )), + )), + vec![0], + true, + ), + ( + "same col, same literal is not contradictory", + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + )), + )), + vec![0], + false, + ), + ]; + + for (desc, expr, expected_cols, expected_infeasible) in cases { + let (result, infeasible) = collect_equality_columns(&expr); + let expected: HashSet<usize> = expected_cols.into_iter().collect(); + if expected_infeasible { + // When infeasible, the scan is short-circuited, so we only + // assert the infeasibility flag — the partial column set + // contents are an implementation detail. + assert!(infeasible, "case '{desc}': expected infeasible"); + } else { + assert_eq!(result, expected, "case '{desc}': columns mismatch"); + assert!(!infeasible, "case '{desc}': expected feasible"); + } + } + } + + /// Regression test: ProjectionExec on top of a FilterExec that already has + /// an explicit projection must not panic when `try_swapping_with_projection` + /// attempts to swap the two nodes. + /// + /// Before the fix, `FilterExecBuilder::from(self)` copied the old projection + /// (e.g. `[0, 1, 2]`) from the FilterExec. After `.with_input` replaced the + /// input with the narrower ProjectionExec (2 columns), `.build()` tried to + /// validate the stale `[0, 1, 2]` projection against the 2-column schema and + /// panicked with "project index 2 out of bounds, max field 2". + #[test] + fn test_filter_with_projection_swap_does_not_panic() -> Result<()> { + use crate::projection::ProjectionExpr; + use datafusion_physical_expr::expressions::col; + + // Schema: [ts: Int64, tokens: Int64, svc: Utf8] + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("tokens", DataType::Int64, false), + Field::new("svc", DataType::Utf8, false), + ])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + // FilterExec: ts > 0, projection=[ts@0, tokens@1, svc@2] (all 3 cols) + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("ts", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int64(Some(0)))), + )); + let filter = Arc::new( + FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![0, 1, 2]))? + .build()?, + ); + + // ProjectionExec: narrows to [ts, tokens] (drops svc) + let proj_exprs = vec![ + ProjectionExpr { + expr: col("ts", &filter.schema())?, + alias: "ts".to_string(), + }, + ProjectionExpr { + expr: col("tokens", &filter.schema())?, + alias: "tokens".to_string(), + }, + ]; + let projection = Arc::new(ProjectionExec::try_new( + proj_exprs, + Arc::clone(&filter) as _, + )?); + + // This must not panic + let result = filter.try_swapping_with_projection(&projection)?; + assert!(result.is_some(), "swap should succeed"); + + let new_plan = result.unwrap(); + // Output schema must still be [ts, tokens] + let out_schema = new_plan.schema(); + assert_eq!(out_schema.fields().len(), 2); + assert_eq!(out_schema.field(0).name(), "ts"); + assert_eq!(out_schema.field(1).name(), "tokens"); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_ndv_capped_at_row_count() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(80), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + // a <= 10 => ~10 rows out of 100 + let predicate: Arc<dyn PhysicalExpr> = + binary(col("a", &schema)?, Operator::LtEq, lit(10i32), &schema)?; + + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + // Filter estimates ~10 rows (selectivity = 10/100) + assert_eq!(statistics.num_rows, Precision::Inexact(10)); + let ndv = &statistics.column_statistics[0].distinct_count; + assert!( + ndv.get_value().copied() <= Some(10), + "Expected NDV <= 10 (filtered row count), got {ndv:?}" + ); + // `a <= 10` rejects nulls, so the 80 input nulls drop to exactly zero. + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + // byte_size follows the same 10% selectivity estimate. + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(100) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_default_selectivity_column_stats() -> Result<()> { + let schema = Schema::new(vec![Field::new("name", DataType::Utf8, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(60), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + // Utf8 interval analysis is unsupported, so this exercises the default + // selectivity path. The predicate rejects nulls but does not constrain + // the column to one value. + let predicate: Arc<dyn PhysicalExpr> = + binary(col("name", &schema)?, Operator::Gt, lit("m"), &schema)?; + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(20)); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(200) + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Inexact(20) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_or_does_not_reject_nulls() -> Result<()> { + let schema = Schema::new(vec![Field::new("name", DataType::Utf8, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(60), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + let predicate: Arc<dyn PhysicalExpr> = binary( + binary(col("name", &schema)?, Operator::Gt, lit("m"), &schema)?, + Operator::Or, + is_null(col("name", &schema)?)?, + &schema, + )?; + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(20)); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Inexact(20) + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(200) + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Inexact(20) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_is_not_null_rejects_nulls() -> Result<()> { + let schema = Schema::new(vec![Field::new("name", DataType::Utf8, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(60), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + // `name IS NOT NULL` keeps only non-null rows, so the surviving null + // count is exactly zero. Utf8 interval analysis is unsupported, so this + // also exercises the default-selectivity path. + let predicate: Arc<dyn PhysicalExpr> = is_not_null(col("name", &schema)?)?; + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(20)); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(200) + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Inexact(20) + ); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/filter_pushdown.rs b/native/vendor/datafusion-physical-plan/src/filter_pushdown.rs new file mode 100644 index 0000000..382967c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/filter_pushdown.rs
@@ -0,0 +1,558 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Filter Pushdown Optimization Process +//! +//! The filter pushdown mechanism involves four key steps: +//! 1. **Optimizer Asks Parent for a Filter Pushdown Plan**: The optimizer calls [`ExecutionPlan::gather_filters_for_pushdown`] +//! on the parent node, passing in parent predicates and phase. The parent node creates a [`FilterDescription`] +//! by inspecting its logic and children's schemas, determining which filters can be pushed to each child. +//! 2. **Optimizer Executes Pushdown**: The optimizer recursively pushes down filters for each child, +//! passing the appropriate filters (`Vec<Arc<dyn PhysicalExpr>>`) for that child. +//! 3. **Optimizer Gathers Results**: The optimizer collects [`FilterPushdownPropagation`] results from children, +//! containing information about which filters were successfully pushed down vs. unsupported. +//! 4. **Parent Responds**: The optimizer calls [`ExecutionPlan::handle_child_pushdown_result`] on the parent, +//! passing a [`ChildPushdownResult`] containing the aggregated pushdown outcomes. The parent decides +//! how to handle filters that couldn't be pushed down (e.g., keep them as FilterExec nodes). +//! +//! [`ExecutionPlan::gather_filters_for_pushdown`]: crate::ExecutionPlan::gather_filters_for_pushdown +//! [`ExecutionPlan::handle_child_pushdown_result`]: crate::ExecutionPlan::handle_child_pushdown_result +//! +//! See also datafusion/physical-optimizer/src/filter_pushdown.rs. + +use std::collections::HashSet; +use std::sync::Arc; + +use arrow_schema::SchemaRef; +use datafusion_common::{ + Result, + tree_node::{Transformed, TreeNode}, +}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FilterPushdownPhase { + /// Pushdown that happens before most other optimizations. + /// This pushdown allows static filters that do not reference any [`ExecutionPlan`]s to be pushed down. + /// Filters that reference an [`ExecutionPlan`] cannot be pushed down at this stage since the whole plan tree may be rewritten + /// by other optimizations. + /// Implementers are however allowed to modify the execution plan themselves during this phase, for example by returning a completely + /// different [`ExecutionPlan`] from [`ExecutionPlan::handle_child_pushdown_result`]. + /// + /// Pushdown of [`FilterExec`] into `DataSourceExec` is an example of a pre-pushdown. + /// Unlike filter pushdown in the logical phase, which operates on the logical plan to push filters into the logical table scan, + /// the `Pre` phase in the physical plan targets the actual physical scan, pushing filters down to specific data source implementations. + /// For example, Parquet supports filter pushdown to reduce data read during scanning, while CSV typically does not. + /// + /// [`ExecutionPlan`]: crate::ExecutionPlan + /// [`FilterExec`]: crate::filter::FilterExec + /// [`ExecutionPlan::handle_child_pushdown_result`]: crate::ExecutionPlan::handle_child_pushdown_result + Pre, + /// Pushdown that happens after most other optimizations. + /// This stage of filter pushdown allows filters that reference an [`ExecutionPlan`] to be pushed down. + /// Since subsequent optimizations should not change the structure of the plan tree except for calling [`ExecutionPlan::with_new_children`] + /// (which generally preserves internal references) it is safe for references between [`ExecutionPlan`]s to be established at this stage. + /// + /// This phase is used to link a [`SortExec`] (with a TopK operator) or a [`HashJoinExec`] to a `DataSourceExec`. + /// + /// [`ExecutionPlan`]: crate::ExecutionPlan + /// [`ExecutionPlan::with_new_children`]: crate::ExecutionPlan::with_new_children + /// [`SortExec`]: crate::sorts::sort::SortExec + /// [`HashJoinExec`]: crate::joins::HashJoinExec + /// [`ExecutionPlan::handle_child_pushdown_result`]: crate::ExecutionPlan::handle_child_pushdown_result + Post, +} + +impl std::fmt::Display for FilterPushdownPhase { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + FilterPushdownPhase::Pre => write!(f, "Pre"), + FilterPushdownPhase::Post => write!(f, "Post"), + } + } +} + +/// The result of a plan for pushing down a filter into a child node. +/// This contains references to filters so that nodes can mutate a filter +/// before pushing it down to a child node (e.g. to adjust a projection) +/// or can directly take ownership of filters that their children +/// could not handle. +#[derive(Debug, Clone)] +pub struct PushedDownPredicate { + pub discriminant: PushedDown, + pub predicate: Arc<dyn PhysicalExpr>, +} + +impl PushedDownPredicate { + /// Return the wrapped [`PhysicalExpr`], discarding whether it is supported or unsupported. + pub fn into_inner(self) -> Arc<dyn PhysicalExpr> { + self.predicate + } + + /// Create a new [`PushedDownPredicate`] with supported pushdown. + pub fn supported(predicate: Arc<dyn PhysicalExpr>) -> Self { + Self { + discriminant: PushedDown::Yes, + predicate, + } + } + + /// Create a new [`PushedDownPredicate`] with unsupported pushdown. + pub fn unsupported(predicate: Arc<dyn PhysicalExpr>) -> Self { + Self { + discriminant: PushedDown::No, + predicate, + } + } +} + +/// Discriminant for the result of pushing down a filter into a child node. +#[derive(Debug, Clone, Copy)] +pub enum PushedDown { + /// The predicate was successfully pushed down into the child node. + Yes, + /// The predicate could not be pushed down into the child node. + No, +} + +impl PushedDown { + /// Logical AND operation: returns `Yes` only if both operands are `Yes`. + pub fn and(self, other: PushedDown) -> PushedDown { + match (self, other) { + (PushedDown::Yes, PushedDown::Yes) => PushedDown::Yes, + _ => PushedDown::No, + } + } + + /// Logical OR operation: returns `Yes` if either operand is `Yes`. + pub fn or(self, other: PushedDown) -> PushedDown { + match (self, other) { + (PushedDown::Yes, _) | (_, PushedDown::Yes) => PushedDown::Yes, + (PushedDown::No, PushedDown::No) => PushedDown::No, + } + } + + /// Wrap a [`PhysicalExpr`] with this pushdown result. + pub fn wrap_expression(self, expr: Arc<dyn PhysicalExpr>) -> PushedDownPredicate { + PushedDownPredicate { + discriminant: self, + predicate: expr, + } + } +} + +/// The result of pushing down a single parent filter into all children. +#[derive(Debug, Clone)] +pub struct ChildFilterPushdownResult { + pub filter: Arc<dyn PhysicalExpr>, + pub child_results: Vec<PushedDown>, +} + +impl ChildFilterPushdownResult { + /// Combine all child results using OR logic. + /// Returns `Yes` if **any** child supports the filter. + /// Returns `No` if **all** children reject the filter or if there are no children. + pub fn any(&self) -> PushedDown { + if self.child_results.is_empty() { + // If there are no children, filters cannot be supported + PushedDown::No + } else { + self.child_results + .iter() + .fold(PushedDown::No, |acc, result| acc.or(*result)) + } + } + + /// Combine all child results using AND logic. + /// Returns `Yes` if **all** children support the filter. + /// Returns `No` if **any** child rejects the filter or if there are no children. + pub fn all(&self) -> PushedDown { + if self.child_results.is_empty() { + // If there are no children, filters cannot be supported + PushedDown::No + } else { + self.child_results + .iter() + .fold(PushedDown::Yes, |acc, result| acc.and(*result)) + } + } +} + +/// The result of pushing down filters into a child node. +/// +/// This is the result provided to nodes in [`ExecutionPlan::handle_child_pushdown_result`]. +/// Nodes process this result and convert it into a [`FilterPushdownPropagation`] +/// that is returned to their parent. +/// +/// [`ExecutionPlan::handle_child_pushdown_result`]: crate::ExecutionPlan::handle_child_pushdown_result +#[derive(Debug, Clone)] +pub struct ChildPushdownResult { + /// The parent filters that were pushed down as received by the current node when [`ExecutionPlan::gather_filters_for_pushdown`](crate::ExecutionPlan::handle_child_pushdown_result) was called. + /// Note that this may *not* be the same as the filters that were passed to the children as the current node may have modified them + /// (e.g. by reassigning column indices) when it returned them from [`ExecutionPlan::gather_filters_for_pushdown`](crate::ExecutionPlan::handle_child_pushdown_result) in a [`FilterDescription`]. + /// Attached to each filter is a [`PushedDown`] *per child* that indicates whether the filter was supported or unsupported by each child. + /// To get combined results see [`ChildFilterPushdownResult::any`] and [`ChildFilterPushdownResult::all`]. + pub parent_filters: Vec<ChildFilterPushdownResult>, + /// The result of pushing down each filter this node provided into each of it's children. + /// The outer vector corresponds to each child, and the inner vector corresponds to each filter. + /// Since this node may have generated a different filter for each child the inner vector may have different lengths or the expressions may not match at all. + /// It is up to each node to interpret this result based on the filters it provided for each child in [`ExecutionPlan::gather_filters_for_pushdown`](crate::ExecutionPlan::handle_child_pushdown_result). + pub self_filters: Vec<Vec<PushedDownPredicate>>, +} + +/// The result of pushing down filters into a node. +/// +/// Returned from [`ExecutionPlan::handle_child_pushdown_result`] to communicate +/// to the optimizer: +/// +/// 1. What to do with any parent filters that could not be pushed down into the children. +/// 2. If the node needs to be replaced in the execution plan with a new node or not. +/// +/// [`ExecutionPlan::handle_child_pushdown_result`]: crate::ExecutionPlan::handle_child_pushdown_result +#[derive(Debug, Clone)] +pub struct FilterPushdownPropagation<T> { + /// Which parent filters were pushed down into this node's children. + pub filters: Vec<PushedDown>, + /// The updated node, if it was updated during pushdown + pub updated_node: Option<T>, +} + +impl<T> FilterPushdownPropagation<T> { + /// Create a new [`FilterPushdownPropagation`] that tells the parent node that each parent filter + /// is supported if it was supported by *all* children. + pub fn if_all(child_pushdown_result: ChildPushdownResult) -> Self { + let filters = child_pushdown_result + .parent_filters + .into_iter() + .map(|result| result.all()) + .collect(); + Self { + filters, + updated_node: None, + } + } + + /// Create a new [`FilterPushdownPropagation`] that tells the parent node that each parent filter + /// is supported if it was supported by *any* child. + pub fn if_any(child_pushdown_result: ChildPushdownResult) -> Self { + let filters = child_pushdown_result + .parent_filters + .into_iter() + .map(|result| result.any()) + .collect(); + Self { + filters, + updated_node: None, + } + } + + /// Create a new [`FilterPushdownPropagation`] that tells the parent node that no filters were pushed down regardless of the child results. + pub fn all_unsupported(child_pushdown_result: ChildPushdownResult) -> Self { + let filters = child_pushdown_result + .parent_filters + .into_iter() + .map(|_| PushedDown::No) + .collect(); + Self { + filters, + updated_node: None, + } + } + + /// Create a new [`FilterPushdownPropagation`] with the specified filter support. + /// This transmits up to our parent node what the result of pushing down the filters into our node and possibly our subtree was. + pub fn with_parent_pushdown_result(filters: Vec<PushedDown>) -> Self { + Self { + filters, + updated_node: None, + } + } + + /// Bind an updated node to the [`FilterPushdownPropagation`]. + /// Use this when the current node wants to update itself in the tree or replace itself with a new node (e.g. one of it's children). + /// You do not need to call this if one of the children of the current node may have updated itself, that is handled by the optimizer. + pub fn with_updated_node(mut self, updated_node: T) -> Self { + self.updated_node = Some(updated_node); + self + } +} + +/// Describes filter pushdown for a single child node. +/// +/// This structure contains two types of filters: +/// - **Parent filters**: Filters received from the parent node, marked as supported or unsupported +/// - **Self filters**: Filters generated by the current node to be pushed down to this child +#[derive(Debug, Clone)] +pub struct ChildFilterDescription { + /// Description of which parent filters can be pushed down into this node. + /// Since we need to transmit filter pushdown results back to this node's parent + /// we need to track each parent filter for each child, even those that are unsupported / won't be pushed down. + /// The entries must stay in the same order as the input parent filters: the + /// filter pushdown optimizer maps child results back to parent filters by + /// position. + pub(crate) parent_filters: Vec<PushedDownPredicate>, + /// Description of which filters this node is pushing down to its children. + /// Since this is not transmitted back to the parents we can have variable sized inner arrays + /// instead of having to track supported/unsupported. + pub(crate) self_filters: Vec<Arc<dyn PhysicalExpr>>, +} + +/// Validates and remaps filter column references to a target schema in one step. +/// +/// When pushing filters from a parent to a child node, we need to: +/// 1. Verify that all columns referenced by the filter exist in the target +/// 2. Remap column indices to match the target schema +/// +/// `allowed_indices` controls which column indices (in the parent schema) are +/// considered valid. For single-input nodes this defaults to +/// `0..child_schema.len()` (all columns are reachable). For join nodes it is +/// restricted to the subset of output columns that map to the target child, +/// which is critical when different sides have same-named columns. +pub(crate) struct FilterRemapper { + /// The target schema to remap column indices into. + child_schema: SchemaRef, + /// Only columns at these indices (in the *parent* schema) are considered + /// valid. For non-join nodes this defaults to `0..child_schema.len()`. + allowed_indices: HashSet<usize>, +} + +impl FilterRemapper { + /// Create a remapper that accepts any column whose index falls within + /// `0..child_schema.len()` and whose name exists in the target schema. + pub(crate) fn new(child_schema: SchemaRef) -> Self { + let allowed_indices = (0..child_schema.fields().len()).collect(); + Self { + child_schema, + allowed_indices, + } + } + + /// Create a remapper that only accepts columns at the given indices. + /// This is used by join nodes to restrict pushdown to one side of the + /// join when both sides have same-named columns. + fn with_allowed_indices( + child_schema: SchemaRef, + allowed_indices: HashSet<usize>, + ) -> Self { + Self { + child_schema, + allowed_indices, + } + } + + /// Try to remap a filter's column references to the target schema. + /// + /// Validates and remaps in a single tree traversal: for each column, + /// checks that its index is in the allowed set and that + /// its name exists in the target schema, then remaps the index. + /// Returns `Some(remapped)` if all columns are valid, or `None` if any + /// column fails validation. + pub(crate) fn try_remap( + &self, + filter: &Arc<dyn PhysicalExpr>, + ) -> Result<Option<Arc<dyn PhysicalExpr>>> { + let mut all_valid = true; + let transformed = Arc::clone(filter).transform_down(|expr| { + if let Some(col) = expr.downcast_ref::<Column>() { + if self.allowed_indices.contains(&col.index()) + && let Ok(new_index) = self.child_schema.index_of(col.name()) + { + Ok(Transformed::yes(Arc::new(Column::new( + col.name(), + new_index, + )))) + } else { + all_valid = false; + Ok(Transformed::complete(expr)) + } + } else { + Ok(Transformed::no(expr)) + } + })?; + + Ok(all_valid.then_some(transformed.data)) + } +} + +impl ChildFilterDescription { + /// Build a child filter description by analyzing which parent filters can be pushed to a specific child. + /// + /// This method performs column analysis to determine which filters can be pushed down: + /// - If all columns referenced by a filter exist in the child's schema, it can be pushed down + /// - Otherwise, it cannot be pushed down to that child + /// + /// See [`FilterDescription::from_children`] for more details + pub fn from_child( + parent_filters: &[Arc<dyn PhysicalExpr>], + child: &Arc<dyn crate::ExecutionPlan>, + ) -> Result<Self> { + let remapper = FilterRemapper::new(child.schema()); + Self::remap_filters(parent_filters, &remapper) + } + + /// Like [`Self::from_child`], but restricts which parent-level columns are + /// considered reachable through this child. + /// + /// `allowed_indices` is the set of column indices (in the *parent* + /// schema) that map to this child's side of a join. A filter is only + /// eligible for pushdown when **every** column index it references + /// appears in `allowed_indices`. + /// + /// This prevents incorrect pushdown when different join sides have + /// columns with the same name: matching on index ensures a filter + /// referencing the right side's `k@2` is not pushed to the left side + /// which also has a column named `k` but at a different index. + pub fn from_child_with_allowed_indices( + parent_filters: &[Arc<dyn PhysicalExpr>], + allowed_indices: HashSet<usize>, + child: &Arc<dyn crate::ExecutionPlan>, + ) -> Result<Self> { + let remapper = + FilterRemapper::with_allowed_indices(child.schema(), allowed_indices); + Self::remap_filters(parent_filters, &remapper) + } + + fn remap_filters( + parent_filters: &[Arc<dyn PhysicalExpr>], + remapper: &FilterRemapper, + ) -> Result<Self> { + let mut child_parent_filters = Vec::with_capacity(parent_filters.len()); + for filter in parent_filters { + if let Some(remapped) = remapper.try_remap(filter)? { + child_parent_filters.push(PushedDownPredicate::supported(remapped)); + } else { + child_parent_filters + .push(PushedDownPredicate::unsupported(Arc::clone(filter))); + } + } + + Ok(Self { + parent_filters: child_parent_filters, + self_filters: vec![], + }) + } + + /// Mark all parent filters as unsupported for this child. + pub fn all_unsupported(parent_filters: &[Arc<dyn PhysicalExpr>]) -> Self { + Self { + parent_filters: parent_filters + .iter() + .map(|f| PushedDownPredicate::unsupported(Arc::clone(f))) + .collect(), + self_filters: vec![], + } + } + + /// Add a self filter (from the current node) to be pushed down to this child. + pub fn with_self_filter(mut self, filter: Arc<dyn PhysicalExpr>) -> Self { + self.self_filters.push(filter); + self + } + + /// Add multiple self filters. + pub fn with_self_filters(mut self, filters: Vec<Arc<dyn PhysicalExpr>>) -> Self { + self.self_filters.extend(filters); + self + } +} + +/// Describes how filters should be pushed down to children. +/// +/// This structure contains filter descriptions for each child node, specifying: +/// - Which parent filters can be pushed down to each child +/// - Which self-generated filters should be pushed down to each child +/// +/// The filter routing is determined by column analysis - filters can only be pushed +/// to children whose schemas contain all the referenced columns. +#[derive(Debug, Clone)] +pub struct FilterDescription { + /// A filter description for each child. + /// This includes which parent filters and which self filters (from the node in question) + /// will get pushed down to each child. + child_filter_descriptions: Vec<ChildFilterDescription>, +} + +impl Default for FilterDescription { + fn default() -> Self { + Self::new() + } +} + +impl FilterDescription { + /// Create a new empty FilterDescription + pub fn new() -> Self { + Self { + child_filter_descriptions: vec![], + } + } + + /// Add a child filter description + pub fn with_child(mut self, child: ChildFilterDescription) -> Self { + self.child_filter_descriptions.push(child); + self + } + + /// Build a filter description by analyzing which parent filters can be pushed to each child. + /// This method automatically determines filter routing based on column analysis: + /// - If all columns referenced by a filter exist in a child's schema, it can be pushed down + /// - Otherwise, it cannot be pushed down to that child + #[expect(clippy::needless_pass_by_value)] + pub fn from_children( + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + children: &[&Arc<dyn crate::ExecutionPlan>], + ) -> Result<Self> { + let mut desc = Self::new(); + + // For each child, create a ChildFilterDescription + for child in children { + desc = desc + .with_child(ChildFilterDescription::from_child(&parent_filters, child)?); + } + + Ok(desc) + } + + /// Mark all parent filters as unsupported for all children. + pub fn all_unsupported( + parent_filters: &[Arc<dyn PhysicalExpr>], + children: &[&Arc<dyn crate::ExecutionPlan>], + ) -> Self { + let mut desc = Self::new(); + for _ in 0..children.len() { + desc = + desc.with_child(ChildFilterDescription::all_unsupported(parent_filters)); + } + desc + } + + pub fn parent_filters(&self) -> Vec<Vec<PushedDownPredicate>> { + self.child_filter_descriptions + .iter() + .map(|d| &d.parent_filters) + .cloned() + .collect() + } + + pub fn self_filters(&self) -> Vec<Vec<Arc<dyn PhysicalExpr>>> { + self.child_filter_descriptions + .iter() + .map(|d| &d.self_filters) + .cloned() + .collect() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/array_map.rs b/native/vendor/datafusion-physical-plan/src/joins/array_map.rs new file mode 100644 index 0000000..4e56cf0 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/array_map.rs
@@ -0,0 +1,601 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_schema::DataType; +use num_traits::AsPrimitive; +use std::mem::size_of; + +use crate::joins::MapOffset; +use crate::joins::chain::traverse_chain; +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; +use arrow::buffer::BooleanBuffer; +use arrow::datatypes::ArrowNumericType; +use datafusion_common::{Result, ScalarValue, internal_err}; + +/// A macro to downcast only supported integer types (up to 64-bit) and invoke a generic function. +/// +/// Usage: `downcast_supported_integer!(data_type => (Method, arg1, arg2, ...))` +/// +/// The `Method` must be an associated method of [`ArrayMap`] that is generic over +/// `<T: ArrowNumericType>` and allow `T::Native: AsPrimitive<u64>`. +macro_rules! downcast_supported_integer { + ($DATA_TYPE:expr => ($METHOD:ident $(, $ARGS:expr)*)) => { + match $DATA_TYPE { + arrow::datatypes::DataType::Int8 => ArrayMap::$METHOD::<arrow::datatypes::Int8Type>($($ARGS),*), + arrow::datatypes::DataType::Int16 => ArrayMap::$METHOD::<arrow::datatypes::Int16Type>($($ARGS),*), + arrow::datatypes::DataType::Int32 => ArrayMap::$METHOD::<arrow::datatypes::Int32Type>($($ARGS),*), + arrow::datatypes::DataType::Int64 => ArrayMap::$METHOD::<arrow::datatypes::Int64Type>($($ARGS),*), + arrow::datatypes::DataType::UInt8 => ArrayMap::$METHOD::<arrow::datatypes::UInt8Type>($($ARGS),*), + arrow::datatypes::DataType::UInt16 => ArrayMap::$METHOD::<arrow::datatypes::UInt16Type>($($ARGS),*), + arrow::datatypes::DataType::UInt32 => ArrayMap::$METHOD::<arrow::datatypes::UInt32Type>($($ARGS),*), + arrow::datatypes::DataType::UInt64 => ArrayMap::$METHOD::<arrow::datatypes::UInt64Type>($($ARGS),*), + _ => { + return internal_err!( + "Unsupported type for ArrayMap: {:?}", + $DATA_TYPE + ); + } + } + }; +} + +/// A dense map for single-column integer join keys within a limited range. +/// +/// Maps join keys to build-side indices using direct array indexing: +/// `data[val - min_val_in_build_side] -> val_idx_in_build_side + 1`. +/// +/// NULL values are ignored on both the build side and the probe side. +/// +/// # Handling Negative Numbers with `wrapping_sub` +/// +/// This implementation supports signed integer ranges (e.g., `[-5, 5]`) efficiently by +/// treating them as `u64` (Two's Complement) and relying on the bitwise properties of +/// wrapping arithmetic (`wrapping_sub`). +/// +/// In Two's Complement representation, `a_signed - b_signed` produces the same bit pattern +/// as `a_unsigned.wrapping_sub(b_unsigned)` (modulo 2^N). This allows us to perform +/// range calculations and zero-based index mapping uniformly for both signed and unsigned +/// types without branching. +/// +/// ## Examples +/// +/// Consider an `Int64` range `[-5, 5]`. +/// * `min_val (-5)` casts to `u64`: `...11111011` (`u64::MAX - 4`) +/// * `max_val (5)` casts to `u64`: `...00000101` (`5`) +/// +/// **1. Range Calculation** +/// +/// ```text +/// In modular arithmetic, this is equivalent to: +/// (5 - (2^64 - 5)) mod 2^64 +/// = (5 - 2^64 + 5) mod 2^64 +/// = (10 - 2^64) mod 2^64 +/// = 10 +/// +/// ``` +/// The resulting `range` (10) correctly represents the size of the interval `[-5, 5]`. +/// +/// **2. Index Lookup (in `get_matched_indices_with_limit_offset`)** +/// +/// For a probe value of `0` (which is stored as `0u64`): +/// ```text +/// In modular arithmetic, this is equivalent to: +/// (0 - (2^64 - 5)) mod 2^64 +/// = (-2^64 + 5) mod 2^64 +/// = 5 +/// ``` +/// This correctly maps `-5` to index `0`, `0` to index `5`, etc. +#[derive(Debug)] +pub struct ArrayMap { + // data[probSideVal-offset] -> valIdxInBuildSide + 1; 0 for absent + data: Vec<u32>, + // min val in buildSide + offset: u64, + // next[buildSideIdx] -> next matching valIdxInBuildSide + 1; 0 for end of chain. + // If next is empty, it means there are no duplicate keys (no conflicts). + // It uses the same chain-based conflict resolution as [`JoinHashMapType`]. + next: Vec<u32>, + num_of_distinct_key: usize, +} + +impl ArrayMap { + pub fn is_supported_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + ) + } + + pub(crate) fn key_to_u64(v: &ScalarValue) -> Option<u64> { + match v { + ScalarValue::Int8(Some(v)) => Some(*v as u64), + ScalarValue::Int16(Some(v)) => Some(*v as u64), + ScalarValue::Int32(Some(v)) => Some(*v as u64), + ScalarValue::Int64(Some(v)) => Some(*v as u64), + ScalarValue::UInt8(Some(v)) => Some(*v as u64), + ScalarValue::UInt16(Some(v)) => Some(*v as u64), + ScalarValue::UInt32(Some(v)) => Some(*v as u64), + ScalarValue::UInt64(Some(v)) => Some(*v), + _ => None, + } + } + + /// Estimates the maximum memory usage for an `ArrayMap` with the given parameters. + /// + pub fn estimate_memory_size(min_val: u64, max_val: u64, num_rows: usize) -> usize { + let range = Self::calculate_range(min_val, max_val); + if range >= usize::MAX as u64 { + return usize::MAX; + } + let size = (range + 1) as usize; + size.saturating_mul(size_of::<u32>()) + .saturating_add(num_rows.saturating_mul(size_of::<u32>())) + } + + pub fn calculate_range(min_val: u64, max_val: u64) -> u64 { + max_val.wrapping_sub(min_val) + } + + #[inline] + fn key_to_index(key: u64, offset: u64, data_len: usize) -> Option<usize> { + let idx = key.wrapping_sub(offset); + if idx < data_len as u64 { + Some(idx as usize) + } else { + None + } + } + + /// Creates a new [`ArrayMap`] from the given array of join keys. + /// + /// Note: This function processes only the non-null values in the input `array`, + /// ignoring any rows where the key is `NULL`. + /// + pub(crate) fn try_new(array: &ArrayRef, min_val: u64, max_val: u64) -> Result<Self> { + let range = Self::calculate_range(min_val, max_val); + if range >= usize::MAX as u64 { + return internal_err!("ArrayMap key range is too large to be allocated."); + } + let size = (range + 1) as usize; + + let mut data: Vec<u32> = vec![0; size]; + let mut next: Vec<u32> = vec![]; + let mut num_of_distinct_key = 0; + + downcast_supported_integer!( + array.data_type() => ( + fill_data, + array, + min_val, + &mut data, + &mut next, + &mut num_of_distinct_key + ) + )?; + + Ok(Self { + data, + offset: min_val, + next, + num_of_distinct_key, + }) + } + + fn fill_data<T: ArrowNumericType>( + array: &ArrayRef, + offset_val: u64, + data: &mut [u32], + next: &mut Vec<u32>, + num_of_distinct_key: &mut usize, + ) -> Result<()> + where + T::Native: AsPrimitive<u64>, + { + let arr = array.as_primitive::<T>(); + // Iterate in reverse to maintain FIFO order when there are duplicate keys. + for (i, val) in arr.iter().enumerate().rev() { + if let Some(val) = val { + let key: u64 = val.as_(); + let Some(idx) = Self::key_to_index(key, offset_val, data.len()) else { + return internal_err!("failed build Array idx >= data.len()"); + }; + + if data[idx] != 0 { + if next.is_empty() { + *next = vec![0; array.len()] + } + next[i] = data[idx] + } else { + *num_of_distinct_key += 1; + } + data[idx] = (i) as u32 + 1; + } + } + Ok(()) + } + + pub fn num_of_distinct_key(&self) -> usize { + self.num_of_distinct_key + } + + /// Returns the memory usage of this [`ArrayMap`] in bytes. + pub fn size(&self) -> usize { + self.data.capacity() * size_of::<u32>() + self.next.capacity() * size_of::<u32>() + } + + pub fn get_matched_indices_with_limit_offset( + &self, + prob_side_keys: &[ArrayRef], + limit: usize, + current_offset: MapOffset, + probe_indices: &mut Vec<u32>, + build_indices: &mut Vec<u64>, + ) -> Result<Option<MapOffset>> { + if prob_side_keys.len() != 1 { + return internal_err!( + "ArrayMap expects 1 join key, but got {}", + prob_side_keys.len() + ); + } + let array = &prob_side_keys[0]; + + downcast_supported_integer!( + array.data_type() => ( + lookup_and_get_indices, + self, + array, + limit, + current_offset, + probe_indices, + build_indices + ) + ) + } + + /// Looks up `key` (a raw probe value cast to `u64`) in the build side, + /// returning the 1-based build-side slot if the key maps to a non-empty + /// bucket, or `None` otherwise. + #[inline] + fn get_value(&self, key: u64) -> Option<u32> { + let idx = Self::key_to_index(key, self.offset, self.data.len())?; + let value = self.data[idx]; + (value != 0).then_some(value) + } + + fn lookup_and_get_indices<T: ArrowNumericType>( + &self, + array: &ArrayRef, + limit: usize, + current_offset: MapOffset, + probe_indices: &mut Vec<u32>, + build_indices: &mut Vec<u64>, + ) -> Result<Option<MapOffset>> + where + T::Native: Copy + AsPrimitive<u64>, + { + probe_indices.clear(); + build_indices.clear(); + + let arr = array.as_primitive::<T>(); + + let have_null = arr.null_count() > 0; + + if self.next.is_empty() { + for prob_idx in current_offset.0..arr.len() { + if build_indices.len() == limit { + return Ok(Some((prob_idx, None))); + } + + // short circuit + if have_null && arr.is_null(prob_idx) { + continue; + } + // SAFETY: prob_idx is guaranteed to be within bounds by the loop range. + let prob_val: u64 = unsafe { arr.value_unchecked(prob_idx) }.as_(); + let Some(build_value) = self.get_value(prob_val) else { + continue; + }; + build_indices.push((build_value - 1) as u64); + probe_indices.push(prob_idx as u32); + } + Ok(None) + } else { + let mut remaining_output = limit; + let to_skip = match current_offset { + // None `initial_next_idx` indicates that `initial_idx` processing hasn't been started + (idx, None) => idx, + // Zero `initial_next_idx` indicates that `initial_idx` has been processed during + // previous iteration, and it should be skipped + (idx, Some(0)) => idx + 1, + // Otherwise, process remaining `initial_idx` matches by traversing `next_chain`, + // to start with the next index + (idx, Some(next_idx)) => { + let is_last = idx == arr.len() - 1; + if let Some(next_offset) = traverse_chain( + &self.next, + idx, + next_idx as u32, + &mut remaining_output, + probe_indices, + build_indices, + is_last, + ) { + return Ok(Some(next_offset)); + } + idx + 1 + } + }; + + for prob_side_idx in to_skip..arr.len() { + if remaining_output == 0 { + return Ok(Some((prob_side_idx, None))); + } + + if have_null && arr.is_null(prob_side_idx) { + continue; + } + + let is_last = prob_side_idx == arr.len() - 1; + + // SAFETY: prob_idx is guaranteed to be within bounds by the loop range. + let prob_val: u64 = unsafe { arr.value_unchecked(prob_side_idx) }.as_(); + let Some(build_idx) = self.get_value(prob_val) else { + continue; + }; + + if let Some(offset) = traverse_chain( + &self.next, + prob_side_idx, + build_idx, + &mut remaining_output, + probe_indices, + build_indices, + is_last, + ) { + return Ok(Some(offset)); + } + } + Ok(None) + } + } + + pub fn contain_keys(&self, probe_side_keys: &[ArrayRef]) -> Result<BooleanArray> { + if probe_side_keys.len() != 1 { + return internal_err!( + "ArrayMap join expects 1 join key, but got {}", + probe_side_keys.len() + ); + } + let array = &probe_side_keys[0]; + + downcast_supported_integer!( + array.data_type() => ( + contain_keys_helper, + self, + array + ) + ) + } + + fn contain_keys_helper<T: ArrowNumericType>( + &self, + array: &ArrayRef, + ) -> Result<BooleanArray> + where + T::Native: AsPrimitive<u64>, + { + let arr = array.as_primitive::<T>(); + let buffer = BooleanBuffer::collect_bool(arr.len(), |i| { + if arr.is_null(i) { + return false; + } + // SAFETY: i is within bounds [0, arr.len()) + let key: u64 = unsafe { arr.value_unchecked(i) }.as_(); + self.get_value(key).is_some() + }); + Ok(BooleanArray::new(buffer, None)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::array::Int64Array; + use arrow::array::UInt64Array; + use std::sync::Arc; + + #[test] + fn test_array_map_limit_offset_duplicate_elements() -> Result<()> { + let build: ArrayRef = Arc::new(Int32Array::from(vec![1, 1, 2])); + let map = ArrayMap::try_new(&build, 1, 2)?; + let probe = [Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef]; + + let mut prob_idx = Vec::new(); + let mut build_idx = Vec::new(); + let mut next = Some((0, None)); + let mut results = vec![]; + + while let Some(o) = next { + next = map.get_matched_indices_with_limit_offset( + &probe, + 1, + o, + &mut prob_idx, + &mut build_idx, + )?; + results.push((prob_idx.clone(), build_idx.clone(), next)); + } + + let expected = vec![ + (vec![0], vec![0], Some((0, Some(2)))), + (vec![0], vec![1], Some((0, Some(0)))), + (vec![1], vec![2], None), + ]; + assert_eq!(results, expected); + Ok(()) + } + + #[test] + fn test_array_map_with_limit_and_misses() -> Result<()> { + let build: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let map = ArrayMap::try_new(&build, 1, 2)?; + let probe = [Arc::new(Int32Array::from(vec![10, 1, 2])) as ArrayRef]; + + let (mut p_idx, mut b_idx) = (vec![], vec![]); + // Skip 10, find 1, next is 2 + let next = map.get_matched_indices_with_limit_offset( + &probe, + 1, + (0, None), + &mut p_idx, + &mut b_idx, + )?; + assert_eq!(p_idx, vec![1]); + assert_eq!(b_idx, vec![0]); + assert_eq!(next, Some((2, None))); + + // Find 2, end + let next = map.get_matched_indices_with_limit_offset( + &probe, + 1, + next.unwrap(), + &mut p_idx, + &mut b_idx, + )?; + assert_eq!(p_idx, vec![2]); + assert_eq!(b_idx, vec![1]); + assert!(next.is_none()); + Ok(()) + } + + #[test] + fn test_array_map_with_build_duplicates_and_misses() -> Result<()> { + let build_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 1])); + let array_map = ArrayMap::try_new(&build_array, 1, 1)?; + // prob: 10(m), 1(h1, h2), 20(m), 1(h1, h2) + let probe_array: ArrayRef = Arc::new(Int32Array::from(vec![10, 1, 20, 1])); + let prob_side_keys = [probe_array]; + + let mut prob_indices = Vec::new(); + let mut build_indices = Vec::new(); + + // batch_size=3, should get 2 matches from first '1' and 1 match from second '1' + let result_offset = array_map.get_matched_indices_with_limit_offset( + &prob_side_keys, + 3, + (0, None), + &mut prob_indices, + &mut build_indices, + )?; + + assert_eq!(prob_indices, vec![1, 1, 3]); + assert_eq!(build_indices, vec![0, 1, 0]); + assert_eq!(result_offset, Some((3, Some(2)))); + Ok(()) + } + + #[test] + fn test_array_map_rejects_large_out_of_range_probe_key() -> Result<()> { + let build: ArrayRef = + Arc::new(UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])); + let map = ArrayMap::try_new(&build, 0, 10)?; + + assert_eq!(ArrayMap::key_to_index(3, 0, 11), Some(3)); + + // Pick a key for which the computed bucket offset is larger than + // u32::MAX but has low 32 bits equal to 3. It must be bounds-checked + // before casting to usize, otherwise 32-bit targets can truncate it + // into range. + let out_of_range_key = (1_u64 << 32) + 3; + assert_eq!(ArrayMap::key_to_index(out_of_range_key, 0, 11), None); + + let probe = [Arc::new(UInt64Array::from(vec![ + Some(3), + Some(out_of_range_key), + Some(11), + None, + ])) as ArrayRef]; + + let mut matched_probe_indices = vec![]; + let mut matched_build_indices = vec![]; + let next = map.get_matched_indices_with_limit_offset( + &probe, + 10, + (0, None), + &mut matched_probe_indices, + &mut matched_build_indices, + )?; + assert_eq!(matched_probe_indices, vec![0]); + assert_eq!(matched_build_indices, vec![3]); + assert!(next.is_none()); + + let contains = map.contain_keys(&probe)?; + assert!(contains.value(0)); + assert!(!contains.value(1)); + assert!(!contains.value(2)); + assert!(!contains.value(3)); + + Ok(()) + } + + #[test] + fn test_array_map_i64_with_negative_and_positive_numbers() -> Result<()> { + // Build array with a mix of negative and positive i64 values, no duplicates + let build_array: ArrayRef = Arc::new(Int64Array::from(vec![-5, 0, 5, -2, 3, 10])); + let min_val = -5_i128; + let max_val = 10_i128; + + let array_map = ArrayMap::try_new(&build_array, min_val as u64, max_val as u64)?; + + // Probe array + let probe_array: ArrayRef = Arc::new(Int64Array::from(vec![0, -5, 10, -1])); + let prob_side_keys = [Arc::clone(&probe_array)]; + + let mut prob_indices = Vec::new(); + let mut build_indices = Vec::new(); + + // Call once to get all matches + let result_offset = array_map.get_matched_indices_with_limit_offset( + &prob_side_keys, + 10, // A batch size larger than number of probes + (0, None), + &mut prob_indices, + &mut build_indices, + )?; + + // Expected matches, in probe-side order: + // Probe 0 (value 0) -> Build 1 (value 0) + // Probe 1 (value -5) -> Build 0 (value -5) + // Probe 2 (value 10) -> Build 5 (value 10) + let expected_prob_indices = vec![0, 1, 2]; + let expected_build_indices = vec![1, 0, 5]; + + assert_eq!(prob_indices, expected_prob_indices); + assert_eq!(build_indices, expected_build_indices); + assert!(result_offset.is_none()); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/chain.rs b/native/vendor/datafusion-physical-plan/src/joins/chain.rs new file mode 100644 index 0000000..846b750 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/chain.rs
@@ -0,0 +1,69 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::Debug; +use std::ops::Sub; + +use arrow::datatypes::ArrowNativeType; + +use crate::joins::MapOffset; + +/// Traverses the chain of matching indices, collecting results up to the remaining limit. +/// Returns `Some(offset)` if the limit was reached and there are more results to process, +/// or `None` if the chain was fully traversed. +#[inline(always)] +pub(crate) fn traverse_chain<T>( + next_chain: &[T], + prob_idx: usize, + start_chain_idx: T, + remaining: &mut usize, + input_indices: &mut Vec<u32>, + match_indices: &mut Vec<u64>, + is_last_input: bool, +) -> Option<MapOffset> +where + T: Copy + TryFrom<usize> + PartialOrd + Into<u64> + Sub<Output = T>, + <T as TryFrom<usize>>::Error: Debug, + T: ArrowNativeType, +{ + let zero = T::usize_as(0); + let one = T::usize_as(1); + let mut match_row_idx = start_chain_idx - one; + + loop { + match_indices.push(match_row_idx.into()); + input_indices.push(prob_idx as u32); + *remaining -= 1; + + let next = next_chain[match_row_idx.into() as usize]; + + if *remaining == 0 { + // Limit reached - return offset for next call + return if is_last_input && next == zero { + // Finished processing the last input row + None + } else { + Some((prob_idx, Some(next.into()))) + }; + } + if next == zero { + // End of chain + return None; + } + match_row_idx = next - one; + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/cross_join.rs b/native/vendor/datafusion-physical-plan/src/joins/cross_join.rs new file mode 100644 index 0000000..8a477c1 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/cross_join.rs
@@ -0,0 +1,1081 @@ +// 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. + +//! Defines the cross join plan for loading the left side of the cross join +//! and producing batches in parallel for the right partitions + +use std::{sync::Arc, task::Poll}; + +use super::utils::{ + BatchSplitter, BatchTransformer, BuildProbeJoinMetrics, NoopBatchTransformer, + OnceAsync, OnceFut, StatefulStreamResult, adjust_right_output_partitioning, + reorder_output_after_swap, +}; +use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::projection::{ + ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, + physical_to_column_exprs, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::EmptyRecordBatchStream; +use crate::{ + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, handle_state, + validate_child_count, +}; + +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow::compute::concat_batches; +use arrow::datatypes::{Fields, Schema, SchemaRef}; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + JoinType, Result, ScalarValue, assert_eq_or_internal_err, internal_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::equivalence::join_equivalence_properties; + +use async_trait::async_trait; +use futures::{Stream, StreamExt, TryStreamExt, ready}; + +/// Data of the left side that is buffered into memory +#[derive(Debug)] +struct JoinLeftData { + /// Single RecordBatch with all rows from the left side + merged_batch: RecordBatch, + /// Track memory reservation for merged_batch. Relies on drop + /// semantics to release reservation when JoinLeftData is dropped. + _reservation: MemoryReservation, +} + +#[expect(rustdoc::private_intra_doc_links)] +/// Cross Join Execution Plan +/// +/// This operator is used when there are no predicates between two tables and +/// returns the Cartesian product of the two tables. +/// +/// Buffers the left input into memory and then streams batches from each +/// partition on the right input combining them with the buffered left input +/// to generate the output. +/// +/// # Clone / Shared State +/// +/// Note this structure includes a [`OnceAsync`] that is used to coordinate the +/// loading of the left side with the processing in each output stream. +/// Therefore it can not be [`Clone`] +#[derive(Debug)] +pub struct CrossJoinExec { + /// left (build) side which gets loaded in memory + pub left: Arc<dyn ExecutionPlan>, + /// right (probe) side which are combined with left side + pub right: Arc<dyn ExecutionPlan>, + /// The schema once the join is applied + schema: SchemaRef, + /// Buffered copy of left (build) side in memory. + /// + /// This structure is *shared* across all output streams. + /// + /// Each output stream waits on the `OnceAsync` to signal the completion of + /// the left side loading. + left_fut: OnceAsync<JoinLeftData>, + /// Execution plan metrics + metrics: ExecutionPlanMetricsSet, + /// Properties such as schema, equivalence properties, ordering, partitioning, etc. + cache: Arc<PlanProperties>, +} + +impl CrossJoinExec { + /// Create a new [CrossJoinExec]. + pub fn new(left: Arc<dyn ExecutionPlan>, right: Arc<dyn ExecutionPlan>) -> Self { + // left then right + let (all_columns, metadata) = { + let left_schema = left.schema(); + let right_schema = right.schema(); + let left_fields = left_schema.fields().iter(); + let right_fields = right_schema.fields().iter(); + + let mut metadata = left_schema.metadata().clone(); + metadata.extend(right_schema.metadata().clone()); + + ( + left_fields.chain(right_fields).cloned().collect::<Fields>(), + metadata, + ) + }; + + let schema = Arc::new(Schema::new(all_columns).with_metadata(metadata)); + let cache = Self::compute_properties(&left, &right, Arc::clone(&schema)).unwrap(); + + CrossJoinExec { + left, + right, + schema, + left_fut: Default::default(), + metrics: ExecutionPlanMetricsSet::default(), + cache: Arc::new(cache), + } + } + + /// left (build) side which gets loaded in memory + pub fn left(&self) -> &Arc<dyn ExecutionPlan> { + &self.left + } + + /// right side which gets combined with left side + pub fn right(&self) -> &Arc<dyn ExecutionPlan> { + &self.right + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + ) -> Result<PlanProperties> { + // Calculate equivalence properties + // TODO: Check equivalence properties of cross join, it may preserve + // ordering in some cases. + let eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &JoinType::Full, + schema, + &[false, false], + None, + &[], + )?; + + // Get output partitioning: + // TODO: Optimize the cross join implementation to generate M * N + // partitions. + let output_partitioning = adjust_right_output_partitioning( + right.output_partitioning(), + left.schema().fields.len(), + )?; + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Final, + boundedness_from_children([left, right]), + )) + } + + /// Returns a new `ExecutionPlan` that computes the same join as this one, + /// with the left and right inputs swapped using the specified + /// `partition_mode`. + /// + /// # Notes: + /// + /// This function should be called BEFORE inserting any repartitioning + /// operators on the join's children. Check [`super::HashJoinExec::swap_inputs`] + /// for more details. + pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> { + let new_join = + CrossJoinExec::new(Arc::clone(&self.right), Arc::clone(&self.left)); + reorder_output_after_swap( + Arc::new(new_join), + &self.left.schema(), + &self.right.schema(), + ) + } +} + +/// Asynchronously collect the result of the left child +async fn load_left_input( + stream: SendableRecordBatchStream, + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, +) -> Result<JoinLeftData> { + let left_schema = stream.schema(); + + // Load all batches and count the rows + let (batches, _metrics, reservation) = stream + .try_fold( + (Vec::new(), metrics, reservation), + |(mut batches, metrics, reservation), batch| async { + let batch_size = batch.get_array_memory_size(); + // Reserve memory for incoming batch + reservation.try_grow(batch_size)?; + // Update metrics + metrics.build_mem_used.add(batch_size); + metrics.build_input_batches.add(1); + metrics.build_input_rows.add(batch.num_rows()); + // Push batch to output + batches.push(batch); + Ok((batches, metrics, reservation)) + }, + ) + .await?; + + let merged_batch = concat_batches(&left_schema, &batches)?; + + Ok(JoinLeftData { + merged_batch, + _reservation: reservation, + }) +} + +impl DisplayAs for CrossJoinExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CrossJoinExec") + } + DisplayFormatType::TreeRender => { + // no extra info to display + Ok(()) + } + } + } +} + +impl ExecutionPlan for CrossJoinExec { + fn name(&self) -> &'static str { + "CrossJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + // CrossJoin has no join conditions or expressions + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + left_fut: Default::default(), + cache: Arc::clone(&self.cache), + schema: Arc::clone(&self.schema), + })) + } + ChildrenPropertiesMode::Recompute => Ok(Arc::new(CrossJoinExec::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + ))), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + let new_exec = CrossJoinExec { + left: Arc::clone(&self.left), + right: Arc::clone(&self.right), + schema: Arc::clone(&self.schema), + left_fut: Default::default(), // reset the build side! + metrics: ExecutionPlanMetricsSet::default(), + cache: Arc::clone(&self.cache), + }; + Ok(Arc::new(new_exec)) + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::UnspecifiedDistribution, + ]) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + assert_eq_or_internal_err!( + self.left.output_partitioning().partition_count(), + 1, + "Invalid CrossJoinExec, the output partition count of the left child must be 1,\ + consider using CoalescePartitionsExec or the EnforceDistribution rule" + ); + + let stream = self.right.execute(partition, Arc::clone(&context))?; + + let join_metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); + + // Initialization of operator-level reservation + let reservation = + MemoryConsumer::new("CrossJoinExec").register(context.memory_pool()); + + let batch_size = context.session_config().batch_size(); + let enforce_batch_size_in_joins = + context.session_config().enforce_batch_size_in_joins(); + + let left_fut = self.left_fut.try_once(|| { + let left_stream = self.left.execute(0, context)?; + + Ok(load_left_input( + left_stream, + join_metrics.clone(), + reservation, + )) + })?; + + if enforce_batch_size_in_joins { + Ok(Box::pin(CrossJoinStream { + schema: Arc::clone(&self.schema), + left_fut, + right: stream, + left_index: 0, + join_metrics, + state: CrossJoinStreamState::WaitBuildSide, + left_data: RecordBatch::new_empty(self.left().schema()), + batch_transformer: BatchSplitter::new(batch_size), + })) + } else { + Ok(Box::pin(CrossJoinStream { + schema: Arc::clone(&self.schema), + left_fut, + right: stream, + left_index: 0, + join_metrics, + state: CrossJoinStreamState::WaitBuildSide, + left_data: RecordBatch::new_empty(self.left().schema()), + batch_transformer: NoopBatchTransformer::new(), + })) + } + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + // Left side is always broadcast, so it always needs overall stats. + // Right side is partitioned, so it needs per-partition stats. + vec![ChildStats::At(None), ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let left_stats = input_stats[0].as_ref().clone(); + let right_stats = input_stats[1].as_ref().clone(); + + Ok(Arc::new(stats_cartesian_product(left_stats, right_stats))) + } + + /// Tries to swap the projection with its input [`CrossJoinExec`]. If it can be done, + /// it returns the new swapped version having the [`CrossJoinExec`] as the top plan. + /// Otherwise, it returns None. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // Convert projected PhysicalExpr's to columns. If not possible, we cannot proceed. + let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) + else { + return Ok(None); + }; + + let (far_right_left_col_ind, far_left_right_col_ind) = join_table_borders( + self.left().schema().fields().len(), + &projection_as_columns, + ); + + if !join_allows_pushdown( + &projection_as_columns, + &self.schema(), + far_right_left_col_ind, + far_left_right_col_ind, + ) { + return Ok(None); + } + + let (new_left, new_right) = new_join_children( + &projection_as_columns, + far_right_left_col_ind, + far_left_right_col_ind, + self.left(), + self.right(), + )?; + + Ok(Some(Arc::new(CrossJoinExec::new( + Arc::new(new_left), + Arc::new(new_right), + )))) + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CrossJoin(Box::new( + protobuf::CrossJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CrossJoinExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let crossjoin = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CrossJoin, + "CrossJoinExec", + ); + + let left = ctx.decode_required_child( + crossjoin.left.as_deref(), + "CrossJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + crossjoin.right.as_deref(), + "CrossJoinExec", + "right", + )?; + + Ok(Arc::new(CrossJoinExec::new(left, right))) + } +} + +/// [left/right]_col_count are required in case the column statistics are None +fn stats_cartesian_product( + left_stats: Statistics, + right_stats: Statistics, +) -> Statistics { + let left_row_count = left_stats.num_rows; + let right_row_count = right_stats.num_rows; + + // Calculate global stats + let num_rows = left_row_count.multiply(&right_row_count); + + // Each output row includes every left and right column, so the left side is + // repeated once per right row and the right side once per left row. + let left_byte_size = left_stats.total_byte_size.multiply(&right_row_count); + let right_byte_size = right_stats.total_byte_size.multiply(&left_row_count); + let total_byte_size = left_byte_size.add(&right_byte_size); + + let left_col_stats = left_stats.column_statistics; + let right_col_stats = right_stats.column_statistics; + + // the null counts must be multiplied by the row counts of the other side (if defined) + // Min, max and distinct_count on the other hand are invariants. + let cross_join_stats = left_col_stats + .into_iter() + .map(|s| { + let widened_sum = s.sum_value.cast_to_sum_type(); + ColumnStatistics { + null_count: s.null_count.multiply(&right_row_count), + distinct_count: s.distinct_count, + min_value: s.min_value, + max_value: s.max_value, + sum_value: widened_sum + .get_value() + // Cast the row count into the same type as any existing sum value + .and_then(|v| { + Precision::<ScalarValue>::from(right_row_count) + .cast_to(&v.data_type()) + .ok() + }) + .map(|row_count| widened_sum.multiply(&row_count)) + .unwrap_or(Precision::Absent), + byte_size: Precision::Absent, + } + }) + .chain(right_col_stats.into_iter().map(|s| { + let widened_sum = s.sum_value.cast_to_sum_type(); + ColumnStatistics { + null_count: s.null_count.multiply(&left_row_count), + distinct_count: s.distinct_count, + min_value: s.min_value, + max_value: s.max_value, + sum_value: widened_sum + .get_value() + // Cast the row count into the same type as any existing sum value + .and_then(|v| { + Precision::<ScalarValue>::from(left_row_count) + .cast_to(&v.data_type()) + .ok() + }) + .map(|row_count| widened_sum.multiply(&row_count)) + .unwrap_or(Precision::Absent), + byte_size: Precision::Absent, + } + })) + .collect(); + + Statistics { + num_rows, + total_byte_size, + column_statistics: cross_join_stats, + } +} + +/// A stream that issues [RecordBatch]es as they arrive from the right of the join. +struct CrossJoinStream<T> { + /// Input schema + schema: Arc<Schema>, + /// Future for data from left side + left_fut: OnceFut<JoinLeftData>, + /// Right side stream + right: SendableRecordBatchStream, + /// Current value on the left + left_index: usize, + /// Join execution metrics + join_metrics: BuildProbeJoinMetrics, + /// State of the stream + state: CrossJoinStreamState, + /// Left data (copy of the entire buffered left side) + left_data: RecordBatch, + /// Batch transformer + batch_transformer: T, +} + +impl<T: BatchTransformer + Unpin + Send> RecordBatchStream for CrossJoinStream<T> { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Represents states of CrossJoinStream +enum CrossJoinStreamState { + WaitBuildSide, + FetchProbeBatch, + /// Holds the currently processed right side batch + BuildBatches(RecordBatch), +} + +impl CrossJoinStreamState { + /// Tries to extract RecordBatch from CrossJoinStreamState enum. + /// Returns an error if state is not BuildBatches state. + fn try_as_record_batch(&mut self) -> Result<&RecordBatch> { + match self { + CrossJoinStreamState::BuildBatches(rb) => Ok(rb), + _ => internal_err!("Expected RecordBatch in BuildBatches state"), + } + } +} + +fn build_batch( + left_index: usize, + batch: &RecordBatch, + left_data: &RecordBatch, + schema: &Schema, +) -> Result<RecordBatch> { + // Repeat value on the left n times + let arrays = left_data + .columns() + .iter() + .map(|arr| { + let scalar = ScalarValue::try_from_array(arr, left_index)?; + scalar.to_array_of_size(batch.num_rows()) + }) + .collect::<Result<Vec<_>>>()?; + + RecordBatch::try_new_with_options( + Arc::new(schema.clone()), + arrays + .iter() + .chain(batch.columns().iter()) + .cloned() + .collect(), + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(Into::into) +} + +#[async_trait] +impl<T: BatchTransformer + Unpin + Send> Stream for CrossJoinStream<T> { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_impl(cx) + } +} + +impl<T: BatchTransformer> CrossJoinStream<T> { + /// Separate implementation function that unpins the [`CrossJoinStream`] so + /// that partial borrows work correctly + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + return match self.state { + CrossJoinStreamState::WaitBuildSide => { + handle_state!(ready!(self.collect_build_side(cx))) + } + CrossJoinStreamState::FetchProbeBatch => { + handle_state!(ready!(self.fetch_probe_batch(cx))) + } + CrossJoinStreamState::BuildBatches(_) => { + let poll = handle_state!(self.build_batches()); + self.join_metrics.baseline.record_poll(poll) + } + }; + } + } + + /// Collects build (left) side of the join into the state. In case of an empty build batch, + /// the execution terminates. Otherwise, the state is updated to fetch probe (right) batch. + fn collect_build_side( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + let build_timer = self.join_metrics.build_time.timer(); + let left_data = match ready!(self.left_fut.get(cx)) { + Ok(left_data) => left_data, + Err(e) => return Poll::Ready(Err(e)), + }; + build_timer.done(); + + let left_data = left_data.merged_batch.clone(); + let result = if left_data.num_rows() == 0 { + StatefulStreamResult::Ready(None) + } else { + self.left_data = left_data; + self.state = CrossJoinStreamState::FetchProbeBatch; + StatefulStreamResult::Continue + }; + Poll::Ready(Ok(result)) + } + + /// Fetches the probe (right) batch, updates the metrics, and save the batch in the state. + /// Then, the state is updated to build result batches. + fn fetch_probe_batch( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + self.left_index = 0; + let right_data = match ready!(self.right.poll_next_unpin(cx)) { + Some(Ok(right_data)) => right_data, + Some(Err(e)) => return Poll::Ready(Err(e)), + None => { + // Release the right (probe) input pipeline's resources. + let right_schema = self.right.schema(); + self.right = Box::pin(EmptyRecordBatchStream::new(right_schema)); + return Poll::Ready(Ok(StatefulStreamResult::Ready(None))); + } + }; + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(right_data.num_rows()); + + self.state = CrossJoinStreamState::BuildBatches(right_data); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Joins the indexed row of left data with the current probe batch. + /// If all the results are produced, the state is set to fetch new probe batch. + fn build_batches(&mut self) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + let right_batch = self.state.try_as_record_batch()?; + if self.left_index < self.left_data.num_rows() { + match self.batch_transformer.next() { + None => { + let join_timer = self.join_metrics.join_time.timer(); + let result = build_batch( + self.left_index, + right_batch, + &self.left_data, + &self.schema, + ); + join_timer.done(); + + self.batch_transformer.set_batch(result?); + } + Some((batch, last)) => { + if last { + self.left_index += 1; + } + + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + } + } else { + self.state = CrossJoinStreamState::FetchProbeBatch; + } + Ok(StatefulStreamResult::Continue) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common; + use crate::test::{assert_join_metrics, build_table_scan_i32}; + + use datafusion_common::{assert_contains, test_util::batches_to_sort_string}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use insta::assert_snapshot; + + async fn join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, + ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> { + let join = CrossJoinExec::new(left, right); + let columns_header = columns(&join.schema()); + + let stream = join.execute(0, context)?; + let batches = common::collect(stream).await?; + let metrics = join.metrics().unwrap(); + + Ok((columns_header, batches, metrics)) + } + + #[tokio::test] + async fn test_stats_cartesian_product() { + let left_row_count = 11; + let left_bytes = 23; + let right_row_count = 7; + let right_bytes = 27; + + let left = Statistics { + num_rows: Precision::Exact(left_row_count), + total_byte_size: Precision::Exact(left_bytes), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(5), + max_value: Precision::Exact(ScalarValue::Int64(Some(21))), + min_value: Precision::Exact(ScalarValue::Int64(Some(-4))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(42))), + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::from("x")), + min_value: Precision::Exact(ScalarValue::from("a")), + sum_value: Precision::Absent, + null_count: Precision::Exact(3), + byte_size: Precision::Absent, + }, + ], + }; + + let right = Statistics { + num_rows: Precision::Exact(right_row_count), + total_byte_size: Precision::Exact(right_bytes), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Exact(3), + max_value: Precision::Exact(ScalarValue::Int64(Some(12))), + min_value: Precision::Exact(ScalarValue::Int64(Some(0))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(20))), + null_count: Precision::Exact(2), + byte_size: Precision::Absent, + }], + }; + + let result = stats_cartesian_product(left, right); + + let expected = Statistics { + num_rows: Precision::Exact(left_row_count * right_row_count), + total_byte_size: Precision::Exact( + left_bytes * right_row_count + right_bytes * left_row_count, + ), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(5), + max_value: Precision::Exact(ScalarValue::Int64(Some(21))), + min_value: Precision::Exact(ScalarValue::Int64(Some(-4))), + sum_value: Precision::Exact(ScalarValue::Int64(Some( + 42 * right_row_count as i64, + ))), + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::from("x")), + min_value: Precision::Exact(ScalarValue::from("a")), + sum_value: Precision::Absent, + null_count: Precision::Exact(3 * right_row_count), + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Exact(3), + max_value: Precision::Exact(ScalarValue::Int64(Some(12))), + min_value: Precision::Exact(ScalarValue::Int64(Some(0))), + sum_value: Precision::Exact(ScalarValue::Int64(Some( + 20 * left_row_count as i64, + ))), + null_count: Precision::Exact(2 * left_row_count), + byte_size: Precision::Absent, + }, + ], + }; + + assert_eq!(result, expected); + } + + #[tokio::test] + async fn test_stats_cartesian_product_with_unknown_size() { + let left_row_count = 11; + + let left = Statistics { + num_rows: Precision::Exact(left_row_count), + total_byte_size: Precision::Exact(23), + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(5), + max_value: Precision::Exact(ScalarValue::Int64(Some(21))), + min_value: Precision::Exact(ScalarValue::Int64(Some(-4))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(42))), + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::from("x")), + min_value: Precision::Exact(ScalarValue::from("a")), + sum_value: Precision::Absent, + null_count: Precision::Exact(3), + byte_size: Precision::Absent, + }, + ], + }; + + let right = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Exact(3), + max_value: Precision::Exact(ScalarValue::Int64(Some(12))), + min_value: Precision::Exact(ScalarValue::Int64(Some(0))), + sum_value: Precision::Exact(ScalarValue::Int64(Some(20))), + null_count: Precision::Exact(2), + byte_size: Precision::Absent, + }], + }; + + let result = stats_cartesian_product(left, right); + + let expected = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![ + ColumnStatistics { + distinct_count: Precision::Exact(5), + max_value: Precision::Exact(ScalarValue::Int64(Some(21))), + min_value: Precision::Exact(ScalarValue::Int64(Some(-4))), + sum_value: Precision::Absent, // we don't know the row count on the right + null_count: Precision::Absent, // we don't know the row count on the right + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::from("x")), + min_value: Precision::Exact(ScalarValue::from("a")), + sum_value: Precision::Absent, + null_count: Precision::Absent, // we don't know the row count on the right + byte_size: Precision::Absent, + }, + ColumnStatistics { + distinct_count: Precision::Exact(3), + max_value: Precision::Exact(ScalarValue::Int64(Some(12))), + min_value: Precision::Exact(ScalarValue::Int64(Some(0))), + sum_value: Precision::Exact(ScalarValue::Int64(Some( + 20 * left_row_count as i64, + ))), + null_count: Precision::Exact(2 * left_row_count), + byte_size: Precision::Absent, + }, + ], + }; + + assert_eq!(result, expected); + } + + #[tokio::test] + async fn test_stats_cartesian_product_unsigned_sum_widens_to_u64() { + let left_row_count = 2; + let right_row_count = 3; + + let left = Statistics { + num_rows: Precision::Exact(left_row_count), + total_byte_size: Precision::Exact(10), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Exact(2), + max_value: Precision::Exact(ScalarValue::UInt32(Some(10))), + min_value: Precision::Exact(ScalarValue::UInt32(Some(1))), + sum_value: Precision::Exact(ScalarValue::UInt32(Some(7))), + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }], + }; + + let right = Statistics { + num_rows: Precision::Exact(right_row_count), + total_byte_size: Precision::Exact(10), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Exact(3), + max_value: Precision::Exact(ScalarValue::UInt32(Some(12))), + min_value: Precision::Exact(ScalarValue::UInt32(Some(0))), + sum_value: Precision::Exact(ScalarValue::UInt32(Some(11))), + null_count: Precision::Exact(0), + byte_size: Precision::Absent, + }], + }; + + let result = stats_cartesian_product(left, right); + + assert_eq!( + result.column_statistics[0].sum_value, + Precision::Exact(ScalarValue::UInt64(Some(21))) + ); + assert_eq!( + result.column_statistics[1].sum_value, + Precision::Exact(ScalarValue::UInt64(Some(22))) + ); + } + + #[tokio::test] + async fn test_join() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + let left = build_table_scan_i32( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 6]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table_scan_i32( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + + let (columns, batches, metrics) = join_collect(left, right, task_ctx).await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 12 | 14 | + | 1 | 4 | 7 | 11 | 13 | 15 | + | 2 | 5 | 8 | 10 | 12 | 14 | + | 2 | 5 | 8 | 11 | 13 | 15 | + | 3 | 6 | 9 | 10 | 12 | 14 | + | 3 | 6 | 9 | 11 | 13 | 15 | + +----+----+----+----+----+----+ + "); + + assert_join_metrics!(metrics, 6); + + Ok(()) + } + + #[tokio::test] + async fn test_overallocation() -> Result<()> { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + let left = build_table_scan_i32( + ("a1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("b1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("c1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ); + let right = build_table_scan_i32( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + + let err = join_collect(left, right, task_ctx).await.unwrap_err(); + + assert_contains!( + err.to_string(), + "Resources exhausted: Additional allocation failed for CrossJoinExec with top memory consumers (across reservations) as:\n CrossJoinExec" + ); + + Ok(()) + } + + /// Returns the column names on the schema + fn columns(schema: &Schema) -> Vec<String> { + schema.fields().iter().map(|f| f.name().clone()).collect() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs new file mode 100644 index 0000000..ec2968c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/exec.rs
@@ -0,0 +1,7374 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashSet; +use std::fmt; +use std::mem::size_of; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock}; +use std::vec; + +use crate::execution_plan::{ + EmissionType, boundedness_from_children, has_same_children_properties, + plan_contains_expression_id, stub_properties, +}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::joins::Map; +use crate::joins::array_map::ArrayMap; +use crate::joins::hash_join::inlist_builder::build_struct_inlist_values; +use crate::joins::hash_join::shared_bounds::{ + ColumnBounds, PartitionBounds, PushdownStrategy, SharedBuildAccumulator, +}; +use crate::joins::hash_join::stream::{ + BuildSide, BuildSideInitialState, HashJoinStream, HashJoinStreamState, +}; +use crate::joins::join_hash_map::{JoinHashMapU32, JoinHashMapU64}; +use crate::joins::key_range_bitmap::KeyRangeBitmap; +use crate::joins::utils::{ + OnceAsync, OnceFut, asymmetric_join_output_partitioning, reorder_output_after_swap, + swap_join_projection, update_hash, +}; +use crate::joins::{JoinOn, JoinOnRef, PartitionMode, SharedBitmapBuilder}; +use crate::metrics::{Count, MetricBuilder, MetricCategory}; +use crate::projection::{ + EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection, + try_pushdown_through_join_with_column_indices, +}; +use crate::repartition::REPARTITION_RANDOM_STATE; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, ExecutionPlanProperties, ReplaceChildrenOptions, + validate_child_count, +}; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, Statistics, + common::can_project, + joins::utils::{ + BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, + build_join_schema, check_join_is_valid, estimate_join_statistics, + need_produce_result_in_final, symmetric_join_output_partitioning, + }, + metrics::{ExecutionPlanMetricsSet, MetricsSet}, +}; + +use arrow::array::{ArrayRef, BooleanBufferBuilder}; +use arrow::compute::concat_batches; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use arrow::util::bit_util; +use arrow_schema::{DataType, Schema}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::memory::{RecordBatchMemoryCounter, estimate_memory_size}; +use datafusion_common::{ + JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, internal_err, + plan_err, project_schema, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_expr::Accumulator; +use datafusion_functions_aggregate_common::min_max::{MaxAccumulator, MinAccumulator}; +use datafusion_physical_expr::equivalence::{ + ProjectionMapping, join_equivalence_properties, +}; +use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; + +use datafusion_common::hash_utils::RandomState; +use datafusion_physical_expr_common::physical_expr::fmt_sql; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::TryStreamExt; +use parking_lot::Mutex; + +use super::partitioned_hash_eval::SeededRandomState; + +/// Hard-coded seed to ensure hash values from the hash join differ from `RepartitionExec`, avoiding collisions. +pub(crate) const HASH_JOIN_SEED: SeededRandomState = + SeededRandomState::with_seed(12210250226015887276); + +const ARRAY_MAP_CREATED_COUNT_METRIC_NAME: &str = "array_map_created_count"; + +#[expect(clippy::too_many_arguments)] +fn try_create_array_map( + bounds: &Option<PartitionBounds>, + schema: &SchemaRef, + batches: &[RecordBatch], + on_left: &[PhysicalExprRef], + reservation: &mut MemoryReservation, + perfect_hash_join_small_build_threshold: usize, + perfect_hash_join_min_key_density: f64, + null_equality: NullEquality, +) -> Result<Option<(ArrayMap, RecordBatch, Vec<ArrayRef>)>> { + if on_left.len() != 1 { + return Ok(None); + } + + if null_equality == NullEquality::NullEqualsNull { + for batch in batches.iter() { + let arrays = evaluate_expressions_to_arrays(on_left, batch)?; + if arrays[0].null_count() > 0 { + return Ok(None); + } + } + } + + let (min_val, max_val) = if let Some(bounds) = bounds { + let (min_val, max_val) = if let Some(cb) = bounds.get_column_bounds(0) { + (cb.min.clone(), cb.max.clone()) + } else { + return Ok(None); + }; + + if min_val.is_null() || max_val.is_null() { + return Ok(None); + } + + if min_val > max_val { + return internal_err!("min_val>max_val"); + } + + if let Some((mi, ma)) = + ArrayMap::key_to_u64(&min_val).zip(ArrayMap::key_to_u64(&max_val)) + { + (mi, ma) + } else { + return Ok(None); + } + } else { + return Ok(None); + }; + + let range = ArrayMap::calculate_range(min_val, max_val); + let num_row: usize = batches.iter().map(|x| x.num_rows()).sum(); + + // TODO: support create ArrayMap<u64> + if num_row >= u32::MAX as usize { + return Ok(None); + } + + // When the key range spans the full integer domain (e.g. i64::MIN to i64::MAX), + // range is u64::MAX and `range + 1` below would overflow. + if range == usize::MAX as u64 { + return Ok(None); + } + + let dense_ratio = (num_row as f64) / ((range + 1) as f64); + + if range >= perfect_hash_join_small_build_threshold as u64 + && dense_ratio <= perfect_hash_join_min_key_density + { + return Ok(None); + } + + let mem_size = ArrayMap::estimate_memory_size(min_val, max_val, num_row); + reservation.try_grow(mem_size)?; + + let batch = concat_batches(schema, batches)?; + let left_values = evaluate_expressions_to_arrays(on_left, &batch)?; + + let array_map = ArrayMap::try_new(&left_values[0], min_val, max_val)?; + + Ok(Some((array_map, batch, left_values))) +} + +/// HashTable and input data for the left (build side) of a join +pub(super) struct JoinLeftData { + /// The hash table with indices into `batch` + /// Arc is used to allow sharing with SharedBuildAccumulator for hash map pushdown + pub(super) map: Arc<Map>, + /// The input rows for the build side + batch: RecordBatch, + /// The build side on expressions values + values: Vec<ArrayRef>, + /// Shared bitmap builder for visited left indices + visited_indices_bitmap: SharedBitmapBuilder, + /// Counter of running probe-threads, potentially + /// able to update `visited_indices_bitmap` + probe_threads_counter: AtomicUsize, + /// We need to keep this field to maintain accurate memory accounting, even though we don't directly use it. + /// Without holding onto this reservation, the recorded memory usage would become inconsistent with actual usage. + /// This could hide potential out-of-memory issues, especially when upstream operators increase their memory consumption. + /// The MemoryReservation ensures proper tracking of memory resources throughout the join operation's lifecycle. + _reservation: MemoryReservation, + /// Bounds computed from the build side for dynamic filter pushdown. + /// If the partition is empty (no rows) this will be None. + /// If the partition has some rows this will be Some with the bounds for each join key column. + pub(super) bounds: Option<PartitionBounds>, + /// Membership testing strategy for filter pushdown + /// Contains either InList values for small build sides or hash table reference for large build sides + pub(super) membership: PushdownStrategy, + /// Shared atomic flag indicating if any probe partition saw data (for null-aware anti joins) + /// This is shared across all probe partitions to provide global knowledge + pub(super) probe_side_non_empty: AtomicBool, + /// Shared atomic flag indicating if any probe partition saw NULL in join keys (for null-aware anti joins) + pub(super) probe_side_has_null: AtomicBool, +} + +impl JoinLeftData { + /// return a reference to the map + pub(super) fn map(&self) -> &Map { + &self.map + } + + /// returns a reference to the build side batch + pub(super) fn batch(&self) -> &RecordBatch { + &self.batch + } + + /// Returns `true` if the build side physically contains rows. + /// + /// This is distinct from [`Self::has_matchable_build_rows`]: a build side + /// can hold rows while its hash map is empty (see that method). + pub(super) fn has_build_rows(&self) -> bool { + self.batch().num_rows() > 0 + } + + /// Returns `true` if the build-side hash map has any matchable entries. + /// + /// Under [`NullEquality::NullEqualsNothing`] build rows whose join key is + /// NULL are omitted from the map, so this can be `false` even when + /// [`Self::has_build_rows`] is `true`. + pub(super) fn has_matchable_build_rows(&self) -> bool { + !self.map().is_empty() + } + + /// returns a reference to the build side expressions values + pub(super) fn values(&self) -> &[ArrayRef] { + &self.values + } + + /// returns a reference to the visited indices bitmap + pub(super) fn visited_indices_bitmap(&self) -> &SharedBitmapBuilder { + &self.visited_indices_bitmap + } + + /// returns a reference to the InList values for filter pushdown + pub(super) fn membership(&self) -> &PushdownStrategy { + &self.membership + } + + /// Decrements the counter of running threads, and returns `true` + /// if caller is the last running thread + pub(super) fn report_probe_completed(&self) -> bool { + self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1 + } +} + +/// Helps to build [`HashJoinExec`]. +/// +/// Builder can be created from an existing [`HashJoinExec`] using [`From::from`]. +/// In this case, all its fields are inherited. If a field that affects the node's +/// properties is modified, they will be automatically recomputed during the build. +/// +/// # Adding setters +/// +/// When adding a new setter, it is necessary to ensure that the `preserve_properties` +/// flag is set to false if modifying the field requires a recomputation of the plan's +/// properties. +/// +pub struct HashJoinExecBuilder { + exec: HashJoinExec, + preserve_properties: bool, +} + +impl HashJoinExecBuilder { + /// Make a new [`HashJoinExecBuilder`]. + pub fn new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + join_type: JoinType, + ) -> Self { + Self { + exec: HashJoinExec { + left, + right, + on, + filter: None, + join_type, + left_fut: Default::default(), + random_state: HASH_JOIN_SEED, + mode: PartitionMode::Auto, + fetch: None, + metrics: ExecutionPlanMetricsSet::new(), + projection: None, + column_indices: vec![], + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + dynamic_filter: None, + // Will be computed at when plan will be built. + cache: stub_properties(), + join_schema: Arc::new(Schema::empty()), + }, + // As `exec` is initialized with stub properties, + // they will be properly computed when plan will be built. + preserve_properties: false, + } + } + + /// Set join type. + pub fn with_type(mut self, join_type: JoinType) -> Self { + self.exec.join_type = join_type; + self.preserve_properties = false; + self + } + + /// Set projection from the vector. + pub fn with_projection(self, projection: Option<Vec<usize>>) -> Self { + self.with_projection_ref(projection.map(Into::into)) + } + + /// Set projection from the shared reference. + pub fn with_projection_ref(mut self, projection: Option<ProjectionRef>) -> Self { + self.exec.projection = projection; + self.preserve_properties = false; + self + } + + /// Set optional filter. + pub fn with_filter(mut self, filter: Option<JoinFilter>) -> Self { + self.exec.filter = filter; + self + } + + /// Set expressions to join on. + pub fn with_on(mut self, on: Vec<(PhysicalExprRef, PhysicalExprRef)>) -> Self { + self.exec.on = on; + self.preserve_properties = false; + self + } + + /// Set partition mode. + pub fn with_partition_mode(mut self, mode: PartitionMode) -> Self { + self.exec.mode = mode; + self.preserve_properties = false; + self + } + + /// Set null equality property. + pub fn with_null_equality(mut self, null_equality: NullEquality) -> Self { + self.exec.null_equality = null_equality; + self + } + + /// Set null aware property. + pub fn with_null_aware(mut self, null_aware: bool) -> Self { + self.exec.null_aware = null_aware; + self + } + + /// Set fetch property. + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.exec.fetch = fetch; + self + } + + /// Require to recompute plan properties. + pub fn recompute_properties(mut self) -> Self { + self.preserve_properties = false; + self + } + + /// Replace children. + pub fn with_new_children( + mut self, + mut children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Self> { + assert_or_internal_err!( + children.len() == 2, + "wrong number of children passed into `HashJoinExecBuilder`" + ); + self.preserve_properties &= has_same_children_properties(&self.exec, &children)?; + self.exec.right = children.swap_remove(1); + self.exec.left = children.swap_remove(0); + Ok(self) + } + + /// Reset runtime state. + pub fn reset_state(mut self) -> Self { + self.exec.left_fut = Default::default(); + self.exec.dynamic_filter = None; + self.exec.metrics = ExecutionPlanMetricsSet::new(); + self + } + + /// Build result as a dyn execution plan. + pub fn build_exec(self) -> Result<Arc<dyn ExecutionPlan>> { + self.build().map(|p| Arc::new(p) as _) + } + + /// Build resulting execution plan. + pub fn build(self) -> Result<HashJoinExec> { + let Self { + exec, + preserve_properties, + } = self; + + // Validate null_aware flag + if exec.null_aware { + let join_type = exec.join_type(); + if !matches!(join_type, JoinType::LeftAnti) { + return plan_err!( + "null_aware can only be true for LeftAnti joins, got {join_type}" + ); + } + let on = exec.on(); + if on.len() != 1 { + return plan_err!( + "null_aware anti join only supports single column join key, got {} columns", + on.len() + ); + } + } + + if preserve_properties { + return Ok(exec); + } + + let HashJoinExec { + left, + right, + on, + filter, + join_type, + left_fut, + random_state, + mode, + metrics, + projection, + null_equality, + null_aware, + dynamic_filter, + fetch, + // Recomputed. + join_schema: _, + column_indices: _, + cache: _, + } = exec; + + let left_schema = left.schema(); + let right_schema = right.schema(); + if on.is_empty() { + return plan_err!("On constraints in HashJoinExec should be non-empty"); + } + + check_join_is_valid(&left_schema, &right_schema, &on)?; + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + + let join_schema = Arc::new(join_schema); + + // Check if the projection is valid. + can_project(&join_schema, projection.as_deref())?; + + let cache = HashJoinExec::compute_properties( + &left, + &right, + &join_schema, + join_type, + &on, + mode, + projection.as_deref(), + )?; + + Ok(HashJoinExec { + left, + right, + on, + filter, + join_type, + join_schema, + left_fut, + random_state, + mode, + metrics, + projection, + column_indices, + null_equality, + null_aware, + cache: Arc::new(cache), + dynamic_filter, + fetch, + }) + } + + fn with_dynamic_filter(mut self, filter: Option<HashJoinExecDynamicFilter>) -> Self { + self.exec.dynamic_filter = filter; + self + } +} + +impl From<&HashJoinExec> for HashJoinExecBuilder { + fn from(exec: &HashJoinExec) -> Self { + Self { + exec: HashJoinExec { + left: Arc::clone(exec.left()), + right: Arc::clone(exec.right()), + on: exec.on.clone(), + filter: exec.filter.clone(), + join_type: exec.join_type, + join_schema: Arc::clone(&exec.join_schema), + left_fut: Arc::clone(&exec.left_fut), + random_state: exec.random_state.clone(), + mode: exec.mode, + metrics: exec.metrics.clone(), + projection: exec.projection.clone(), + column_indices: exec.column_indices.clone(), + null_equality: exec.null_equality, + null_aware: exec.null_aware, + cache: Arc::clone(&exec.cache), + dynamic_filter: exec.dynamic_filter.clone(), + fetch: exec.fetch, + }, + preserve_properties: true, + } + } +} + +#[expect(rustdoc::private_intra_doc_links)] +/// Join execution plan: Evaluates equijoin predicates in parallel on multiple +/// partitions using a hash table and an optional filter list to apply post +/// join. +/// +/// # Join Expressions +/// +/// This implementation is optimized for evaluating equijoin predicates ( +/// `<col1> = <col2>`) expressions, which are represented as a list of `Columns` +/// in [`Self::on`]. +/// +/// Non-equality predicates, which can not pushed down to a join inputs (e.g. +/// `<col1> != <col2>`) are known as "filter expressions" and are evaluated +/// after the equijoin predicates. +/// +/// # ArrayMap Optimization +/// +/// For joins with a single integer-based join key, `HashJoinExec` may use an [`ArrayMap`] +/// (also known as a "perfect hash join") instead of a general-purpose hash map. +/// This optimization is used when: +/// 1. There is exactly one join key. +/// 2. The join key is an integer type up to 64 bits wide that can be losslessly converted +/// to `u64` (128-bit integer types such as `i128` and `u128` are not supported). +/// 3. The range of keys is small enough (controlled by `perfect_hash_join_small_build_threshold`) +/// OR the keys are sufficiently dense (controlled by `perfect_hash_join_min_key_density`). +/// 4. build_side.num_rows() < u32::MAX +/// 5. NullEqualsNothing || (NullEqualsNull && build side doesn't contain null) +/// +/// See [`try_create_array_map`] for more details. +/// +/// Note that when using [`PartitionMode::Partitioned`], the build side is split into multiple +/// partitions. This can cause a dense build side to become sparse within each partition, +/// potentially disabling this optimization. +/// +/// For example, consider: +/// ```sql +/// SELECT t1.value, t2.value +/// FROM range(10000) AS t1 +/// JOIN range(10000) AS t2 +/// ON t1.value = t2.value; +/// ``` +/// With 24 partitions, each partition will only receive a subset of the 10,000 rows. +/// The first partition might contain values like `3, 10, 18, 39, 43`, which are sparse +/// relative to the original range, even though the overall data set is dense. +/// +/// # "Build Side" vs "Probe Side" +/// +/// HashJoin takes two inputs, which are referred to as the "build" and the +/// "probe". The build side is the first child, and the probe side is the second +/// child. +/// +/// The two inputs are treated differently and it is VERY important that the +/// *smaller* input is placed on the build side to minimize the work of creating +/// the hash table. +/// +/// ```text +/// ┌───────────┐ +/// │ HashJoin │ +/// │ │ +/// └───────────┘ +/// │ │ +/// ┌─────┘ └─────┐ +/// ▼ ▼ +/// ┌────────────┐ ┌─────────────┐ +/// │ Input │ │ Input │ +/// │ [0] │ │ [1] │ +/// └────────────┘ └─────────────┘ +/// +/// "build side" "probe side" +/// ``` +/// +/// Execution proceeds in 2 stages: +/// +/// 1. the **build phase** creates a hash table from the tuples of the build side, +/// and single concatenated batch containing data from all fetched record batches. +/// Resulting hash table stores hashed join-key fields for each row as a key, and +/// indices of corresponding rows in concatenated batch. +/// +/// When using the standard `JoinHashMap`, hash join uses LIFO data structure as a hash table, +/// and in order to retain original build-side input order while obtaining data during probe phase, +/// hash table is updated by iterating batch sequence in reverse order -- it allows to +/// keep rows with smaller indices "on the top" of hash table, and still maintain +/// correct indexing for concatenated build-side data batch. +/// +/// Example of build phase for 3 record batches: +/// +/// +/// ```text +/// +/// Original build-side data Inserting build-side values into hashmap Concatenated build-side batch +/// ┌───────────────────────────┐ +/// hashmap.insert(row-hash, row-idx + offset) │ idx │ +/// ┌───────┐ │ ┌───────┐ │ +/// │ Row 1 │ 1) update_hash for batch 3 with offset 0 │ │ Row 6 │ 0 │ +/// Batch 1 │ │ - hashmap.insert(Row 7, idx 1) │ Batch 3 │ │ │ +/// │ Row 2 │ - hashmap.insert(Row 6, idx 0) │ │ Row 7 │ 1 │ +/// └───────┘ │ └───────┘ │ +/// │ │ +/// ┌───────┐ │ ┌───────┐ │ +/// │ Row 3 │ 2) update_hash for batch 2 with offset 2 │ │ Row 3 │ 2 │ +/// │ │ - hashmap.insert(Row 5, idx 4) │ │ │ │ +/// Batch 2 │ Row 4 │ - hashmap.insert(Row 4, idx 3) │ Batch 2 │ Row 4 │ 3 │ +/// │ │ - hashmap.insert(Row 3, idx 2) │ │ │ │ +/// │ Row 5 │ │ │ Row 5 │ 4 │ +/// └───────┘ │ └───────┘ │ +/// │ │ +/// ┌───────┐ │ ┌───────┐ │ +/// │ Row 6 │ 3) update_hash for batch 1 with offset 5 │ │ Row 1 │ 5 │ +/// Batch 3 │ │ - hashmap.insert(Row 2, idx 6) │ Batch 1 │ │ │ +/// │ Row 7 │ - hashmap.insert(Row 1, idx 5) │ │ Row 2 │ 6 │ +/// └───────┘ │ └───────┘ │ +/// │ │ +/// └───────────────────────────┘ +/// ``` +/// +/// 2. the **probe phase** where the tuples of the probe side are streamed +/// through, checking for matches of the join keys in the hash table. +/// +/// ```text +/// ┌────────────────┐ ┌────────────────┐ +/// │ ┌─────────┐ │ │ ┌─────────┐ │ +/// │ │ Hash │ │ │ │ Hash │ │ +/// │ │ Table │ │ │ │ Table │ │ +/// │ │(keys are│ │ │ │(keys are│ │ +/// │ │equi join│ │ │ │equi join│ │ Stage 2: batches from +/// Stage 1: the │ │columns) │ │ │ │columns) │ │ the probe side are +/// *entire* build │ │ │ │ │ │ │ │ streamed through, and +/// side is read │ └─────────┘ │ │ └─────────┘ │ checked against the +/// into the hash │ ▲ │ │ ▲ │ contents of the hash +/// table │ HashJoin │ │ HashJoin │ table +/// └──────┼─────────┘ └──────────┼─────┘ +/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ +/// │ │ +/// +/// │ │ +/// ┌────────────┐ ┌────────────┐ +/// │RecordBatch │ │RecordBatch │ +/// └────────────┘ └────────────┘ +/// ┌────────────┐ ┌────────────┐ +/// │RecordBatch │ │RecordBatch │ +/// └────────────┘ └────────────┘ +/// ... ... +/// ┌────────────┐ ┌────────────┐ +/// │RecordBatch │ │RecordBatch │ +/// └────────────┘ └────────────┘ +/// +/// build side probe side +/// ``` +/// +/// # Example "Optimal" Plans +/// +/// The differences in the inputs means that for classic "Star Schema Query", +/// the optimal plan will be a **"Right Deep Tree"** . A Star Schema Query is +/// one where there is one large table and several smaller "dimension" tables, +/// joined on `Foreign Key = Primary Key` predicates. +/// +/// A "Right Deep Tree" looks like this large table as the probe side on the +/// lowest join: +/// +/// ```text +/// ┌───────────┐ +/// │ HashJoin │ +/// │ │ +/// └───────────┘ +/// │ │ +/// ┌───────┘ └──────────┐ +/// ▼ ▼ +/// ┌───────────────┐ ┌───────────┐ +/// │ small table 1 │ │ HashJoin │ +/// │ "dimension" │ │ │ +/// └───────────────┘ └───┬───┬───┘ +/// ┌──────────┘ └───────┐ +/// │ │ +/// ▼ ▼ +/// ┌───────────────┐ ┌───────────┐ +/// │ small table 2 │ │ HashJoin │ +/// │ "dimension" │ │ │ +/// └───────────────┘ └───┬───┬───┘ +/// ┌────────┘ └────────┐ +/// │ │ +/// ▼ ▼ +/// ┌───────────────┐ ┌───────────────┐ +/// │ small table 3 │ │ large table │ +/// │ "dimension" │ │ "fact" │ +/// └───────────────┘ └───────────────┘ +/// ``` +/// +/// # Clone / Shared State +/// +/// Note this structure includes a [`OnceAsync`] that is used to coordinate the +/// loading of the left side with the processing in each output stream. +/// Therefore it can not be [`Clone`] +pub struct HashJoinExec { + /// left (build) side which gets hashed + pub left: Arc<dyn ExecutionPlan>, + /// right (probe) side which are filtered by the hash table + pub right: Arc<dyn ExecutionPlan>, + /// Set of equijoin columns from the relations: `(left_col, right_col)` + pub on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + /// Filters which are applied while finding matching rows + pub filter: Option<JoinFilter>, + /// How the join is performed (`OUTER`, `INNER`, etc) + pub join_type: JoinType, + /// The schema after join. Please be careful when using this schema, + /// if there is a projection, the schema isn't the same as the output schema. + join_schema: SchemaRef, + /// Future that consumes left input and builds the hash table + /// + /// For CollectLeft partition mode, this structure is *shared* across all output streams. + /// + /// Each output stream waits on the `OnceAsync` to signal the completion of + /// the hash table creation. + left_fut: Arc<OnceAsync<JoinLeftData>>, + /// Shared the `SeededRandomState` for the hashing algorithm (seeds preserved for serialization) + random_state: SeededRandomState, + /// Partitioning mode to use + pub mode: PartitionMode, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// The projection indices of the columns in the output schema of join + pub projection: Option<ProjectionRef>, + /// Information of index and left / right placement of columns + column_indices: Vec<ColumnIndex>, + /// The equality null-handling behavior of the join algorithm. + pub null_equality: NullEquality, + /// Flag to indicate if this is a null-aware anti join + pub null_aware: bool, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// Dynamic filter for pushing down to the probe side + /// Set when dynamic filter pushdown is detected in handle_child_pushdown_result. + /// HashJoinExec also needs to keep a shared bounds accumulator for coordinating updates. + dynamic_filter: Option<HashJoinExecDynamicFilter>, + /// Maximum number of rows to return + fetch: Option<usize>, +} + +#[derive(Clone)] +struct HashJoinExecDynamicFilter { + /// Dynamic filter that we'll update with the results of the build side once that is done. + filter: Arc<DynamicFilterPhysicalExpr>, + /// Build accumulator to collect build-side information (hash maps and/or bounds) from each partition. + /// It is lazily initialized during execution to make sure we use the actual execution time partition counts. + build_accumulator: OnceLock<Arc<SharedBuildAccumulator>>, +} + +impl fmt::Debug for HashJoinExec { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HashJoinExec") + .field("left", &self.left) + .field("right", &self.right) + .field("on", &self.on) + .field("filter", &self.filter) + .field("join_type", &self.join_type) + .field("join_schema", &self.join_schema) + .field("left_fut", &self.left_fut) + .field("random_state", &self.random_state) + .field("mode", &self.mode) + .field("metrics", &self.metrics) + .field("projection", &self.projection) + .field("column_indices", &self.column_indices) + .field("null_equality", &self.null_equality) + .field("cache", &self.cache) + // Explicitly exclude dynamic_filter to avoid runtime state differences in tests + .finish() + } +} + +impl EmbeddedProjection for HashJoinExec { + fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> { + self.with_projection(projection) + } +} + +impl HashJoinExec { + /// Tries to create a new [`HashJoinExec`]. + /// + /// # Error + /// This function errors when it is not possible to join the left and right sides on keys `on`. + #[expect(clippy::too_many_arguments)] + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: Option<JoinFilter>, + join_type: &JoinType, + projection: Option<Vec<usize>>, + partition_mode: PartitionMode, + null_equality: NullEquality, + null_aware: bool, + ) -> Result<Self> { + HashJoinExecBuilder::new(left, right, on, *join_type) + .with_filter(filter) + .with_projection(projection) + .with_partition_mode(partition_mode) + .with_null_equality(null_equality) + .with_null_aware(null_aware) + .build() + } + + /// Create a builder based on the existing [`HashJoinExec`]. + /// + /// Returned builder preserves all existing fields. If a field requiring properties + /// recomputation is modified, this will be done automatically during the node build. + /// + pub fn builder(&self) -> HashJoinExecBuilder { + self.into() + } + + fn create_dynamic_filter(on: &JoinOn) -> Arc<DynamicFilterPhysicalExpr> { + // Extract the right-side keys (probe side keys) from the `on` clauses + // Dynamic filter will be created from build side values (left side) and applied to probe side (right side) + let right_keys: Vec<_> = on.iter().map(|(_, r)| Arc::clone(r)).collect(); + // Initialize with a placeholder expression (true) that will be updated when the hash table is built + Arc::new(DynamicFilterPhysicalExpr::new(right_keys, lit(true))) + } + + fn allow_join_dynamic_filter_pushdown(&self, config: &ConfigOptions) -> bool { + let (_, probe_preserved) = self.join_type.on_lr_is_preserved(); + if !probe_preserved || !config.optimizer.enable_join_dynamic_filter_pushdown { + return false; + } + + // A null-aware anti join emits a build-side NULL only when the probe + // is truly empty. The pushed filter can empty the probe by pruning + // every row, which would surface that NULL wrongly. A NOT NULL build + // key cannot produce such a NULL, so the filter stays there. + if self.null_aware + && self.on.iter().any(|(build_key, _)| { + build_key.nullable(&self.left.schema()).unwrap_or(true) + }) + { + return false; + } + + // `preserve_file_partitions` can report Hive-style file groups as Hash + // partitioned even though their partition indexes do not follow the + // hash router used by partitioned dynamic filters. Reject Hash inputs + // because the metadata cannot distinguish those scans from a real hash + // repartition. Compatible Range inputs remain safe because matching + // ordering and split points align each build filter with its probe + // partition. Other unsupported layouts are rejected. + // Follow-up work: enable dynamic filtering for preserve_file_partitioned scans (issue #20195). + // https://github.com/apache/datafusion/issues/20195 + if config.optimizer.preserve_file_partitions > 0 + && self.mode == PartitionMode::Partitioned + && matches!( + ( + self.left.output_partitioning(), + self.right.output_partitioning() + ), + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) + ) + { + return false; + } + + if self.mode == PartitionMode::Partitioned + && !self.has_partitioned_dynamic_filter_routing() + { + return false; + } + + true + } + + fn has_partitioned_dynamic_filter_routing(&self) -> bool { + match ( + self.left.output_partitioning(), + self.right.output_partitioning(), + ) { + ( + Partitioning::Hash(_, left_partition_count), + Partitioning::Hash(_, right_partition_count), + ) => left_partition_count == right_partition_count, + (Partitioning::Range(_), Partitioning::Range(_)) => { + let children = [self.left.as_ref(), self.right.as_ref()]; + matches!( + self.input_distribution_requirements() + .unsatisfied_co_partitioned_children(self.name(), &children), + Ok(unsatisfied) if unsatisfied.is_empty() + ) + } + (left_partitioning, right_partitioning) => { + left_partitioning.partition_count() == 1 + && right_partitioning.partition_count() == 1 + } + } + } + + /// left (build) side which gets hashed + pub fn left(&self) -> &Arc<dyn ExecutionPlan> { + &self.left + } + + /// right (probe) side which are filtered by the hash table + pub fn right(&self) -> &Arc<dyn ExecutionPlan> { + &self.right + } + + /// Set of common columns used to join on + pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] { + &self.on + } + + /// Filters applied before join output + pub fn filter(&self) -> Option<&JoinFilter> { + self.filter.as_ref() + } + + /// How the join is performed + pub fn join_type(&self) -> &JoinType { + &self.join_type + } + + /// The schema after join. Please be careful when using this schema, + /// if there is a projection, the schema isn't the same as the output schema. + pub fn join_schema(&self) -> &SchemaRef { + &self.join_schema + } + + /// The partitioning mode of this hash join + pub fn partition_mode(&self) -> &PartitionMode { + &self.mode + } + + /// Get null_equality + pub fn null_equality(&self) -> NullEquality { + self.null_equality + } + + /// Returns the dynamic filter expression produced by this hash join, if set. + #[deprecated( + since = "55.0.0", + note = "Use ExecutionPlan::dynamic_expressions_produced instead" + )] + pub fn dynamic_filter_expr(&self) -> Option<&Arc<DynamicFilterPhysicalExpr>> { + self.dynamic_filter.as_ref().map(|df| &df.filter) + } + + /// Set the dynamic filter on this hash join. + /// + /// Resets any internal state that depends on any existing dynamic filter. + /// + /// Validates that the filter's children reference valid columns in + /// the probe (right) side's schema. + pub fn with_dynamic_filter_expr( + mut self, + filter: Arc<DynamicFilterPhysicalExpr>, + ) -> Result<Self> { + let probe_schema = self.right.schema(); + for child in filter.children() { + child.data_type(&probe_schema)?; + } + self.dynamic_filter = Some(HashJoinExecDynamicFilter { + filter, + // Initialize with an empty accumulator which will be lazily populated + // during execution. + build_accumulator: OnceLock::new(), + }); + Ok(self) + } + + /// Calculate order preservation flags for this hash join. + fn maintains_input_order(join_type: JoinType) -> Vec<bool> { + vec![ + false, + matches!( + join_type, + JoinType::Inner + | JoinType::Right + | JoinType::RightAnti + | JoinType::RightSemi + | JoinType::RightMark + ), + ] + } + + /// Get probe side information for the hash join. + pub fn probe_side() -> JoinSide { + // In current implementation right side is always probe side. + JoinSide::Right + } + + /// Return whether the join contains a projection + pub fn contains_projection(&self) -> bool { + self.projection.is_some() + } + + /// Return new instance of [HashJoinExec] with the given projection. + pub fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> { + let projection = projection.map(Into::into); + // check if the projection is valid + can_project(&self.schema(), projection.as_deref())?; + let projection = + combine_projections(projection.as_ref(), self.projection.as_ref())?; + self.builder().with_projection_ref(projection).build() + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + schema: &SchemaRef, + join_type: JoinType, + on: JoinOnRef, + mode: PartitionMode, + projection: Option<&[usize]>, + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let mut eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &join_type, + Arc::clone(schema), + &Self::maintains_input_order(join_type), + Some(Self::probe_side()), + on, + )?; + + let mut output_partitioning = match mode { + PartitionMode::CollectLeft => { + asymmetric_join_output_partitioning(left, right, &join_type)? + } + PartitionMode::Auto => Partitioning::UnknownPartitioning( + right.output_partitioning().partition_count(), + ), + PartitionMode::Partitioned => { + symmetric_join_output_partitioning(left, right, &join_type)? + } + }; + + let emission_type = if left.boundedness().is_unbounded() { + EmissionType::Final + } else if right.pipeline_behavior() == EmissionType::Incremental { + match join_type { + // If we only need to generate matched rows from the probe side, + // we can emit rows incrementally. + JoinType::Inner + | JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::Right + | JoinType::RightAnti + | JoinType::RightMark => EmissionType::Incremental, + // If we need to generate unmatched rows from the *build side*, + // we need to emit them at the end. + JoinType::Left + | JoinType::LeftAnti + | JoinType::LeftMark + | JoinType::Full => EmissionType::Both, + } + } else { + right.pipeline_behavior() + }; + + // If contains projection, update the PlanProperties. + if let Some(projection) = projection { + // construct a map from the input expressions to the output expression of the Projection + let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; + let out_schema = project_schema(schema, Some(&projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, out_schema); + } + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + emission_type, + boundedness_from_children([left, right]), + )) + } + + /// Returns a new `ExecutionPlan` that computes the same join as this one, + /// with the left and right inputs swapped using the specified + /// `partition_mode`. + /// + /// # Notes: + /// + /// This function is public so other downstream projects can use it to + /// construct `HashJoinExec` with right side as the build side. + /// + /// For using this interface directly, please refer to below: + /// + /// Hash join execution may require specific input partitioning (for example, + /// the left child may have a single partition while the right child has multiple). + /// + /// Calling this function on join nodes whose children have already been repartitioned + /// (e.g., after a `RepartitionExec` has been inserted) may break the partitioning + /// requirements of the hash join. Therefore, ensure you call this function + /// before inserting any repartitioning operators on the join's children. + /// + /// In DataFusion's default SQL interface, this function is used by the `JoinSelection` + /// physical optimizer rule to determine a good join order, which is + /// executed before the `EnforceDistribution` rule (the rule that may + /// insert `RepartitionExec` operators). + pub fn swap_inputs( + &self, + partition_mode: PartitionMode, + ) -> Result<Arc<dyn ExecutionPlan>> { + assert_or_internal_err!( + self.dynamic_filter.is_none(), + "Cannot swap HashJoinExec inputs after dynamic filters have been constructed. \ + Optimizer rules that reorder join inputs must run before optimizer rules `FilterPushdown::new_post_optimization()`" + ); + + let left = self.left(); + let right = self.right(); + let new_join = self + .builder() + .with_type(self.join_type.swap()) + .with_new_children(vec![Arc::clone(right), Arc::clone(left)])? + .with_on( + self.on() + .iter() + .map(|(l, r)| (Arc::clone(r), Arc::clone(l))) + .collect(), + ) + .with_filter(self.filter().map(JoinFilter::swap)) + .with_projection(swap_join_projection( + left.schema().fields().len(), + right.schema().fields().len(), + self.projection.as_deref(), + self.join_type(), + )) + .with_partition_mode(partition_mode) + .build()?; + // In case of anti / semi joins or if there is embedded projection in HashJoinExec, output column order is preserved, no need to add projection again + if matches!( + self.join_type(), + JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ) || self.projection.is_some() + { + Ok(Arc::new(new_join)) + } else { + reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) + } + } +} + +impl DisplayAs for HashJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let display_filter = self.filter.as_ref().map_or_else( + || "".to_string(), + |f| format!(", filter={}", f.expression()), + ); + let display_projections = if self.contains_projection() { + format!( + ", projection=[{}]", + self.projection + .as_ref() + .unwrap() + .iter() + .map(|index| format!( + "{}@{}", + self.join_schema.fields().get(*index).unwrap().name(), + index + )) + .collect::<Vec<_>>() + .join(", ") + ) + } else { + "".to_string() + }; + let display_null_equality = + if self.null_equality() == NullEquality::NullEqualsNull { + ", NullsEqual: true" + } else { + "" + }; + let display_fetch = self + .fetch + .map_or_else(String::new, |f| format!(", fetch={f}")); + let display_null_aware = + if self.null_aware { ", null_aware" } else { "" }; + let on = self + .on + .iter() + .map(|(c1, c2)| format!("({c1}, {c2})")) + .collect::<Vec<String>>() + .join(", "); + write!( + f, + "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}", + self.mode, + self.join_type, + on, + display_filter, + display_projections, + display_null_equality, + display_fetch, + display_null_aware, + ) + } + DisplayFormatType::TreeRender => { + let on = self + .on + .iter() + .map(|(c1, c2)| { + format!("({} = {})", fmt_sql(c1.as_ref()), fmt_sql(c2.as_ref())) + }) + .collect::<Vec<String>>() + .join(", "); + + if *self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + + writeln!(f, "on={on}")?; + + if self.null_equality() == NullEquality::NullEqualsNull { + writeln!(f, "NullsEqual: true")?; + } + + if self.null_aware { + writeln!(f, "null_aware")?; + } + + if let Some(filter) = self.filter.as_ref() { + writeln!(f, "filter={filter}")?; + } + + if let Some(fetch) = self.fetch { + writeln!(f, "fetch={fetch}")?; + } + + Ok(()) + } + } + } +} + +impl ExecutionPlan for HashJoinExec { + fn name(&self) -> &'static str { + "HashJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + match self.mode { + PartitionMode::Partitioned => { + let (left_expr, right_expr) = self + .on + .iter() + .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) + .unzip(); + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) + } + PartitionMode::CollectLeft => InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::UnspecifiedDistribution, + ]), + PartitionMode::Auto => InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + Distribution::UnspecifiedDistribution, + ]), + } + } + + // For [JoinType::Inner] and [JoinType::RightSemi] in hash joins, the probe phase initiates by + // applying the hash function to convert the join key(s) in each row into a hash value from the + // probe side table in the order they're arranged. The hash value is used to look up corresponding + // entries in the hash table that was constructed from the build side table during the build phase. + // + // Because of the immediate generation of result rows once a match is found, + // the output of the join tends to follow the order in which the rows were read from + // the probe side table. This is simply due to the sequence in which the rows were processed. + // Hence, it appears that the hash join is preserving the order of the probe side. + // + // Meanwhile, in the case of a [JoinType::RightAnti] hash join, + // the unmatched rows from the probe side are also kept in order. + // This is because the **`RightAnti`** join is designed to return rows from the right + // (probe side) table that have no match in the left (build side) table. Because the rows + // are processed sequentially in the probe phase, and unmatched rows are directly output + // as results, these results tend to retain the order of the probe side table. + fn maintains_input_order(&self) -> Vec<bool> { + Self::maintains_input_order(self.join_type) + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let join_keys = self + .on + .iter() + .flat_map(|(left, right)| [Arc::clone(left), Arc::clone(right)]); + let filter = self + .filter + .iter() + .map(|filter| Arc::clone(filter.expression())); + let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| { + Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter.filter) + as Arc<dyn PhysicalExpr> + }); + crate::apply_expression_roots(join_keys.chain(filter).chain(dynamic_filter), f) + } + + fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.dynamic_filter + .iter() + .map(|dynamic_filter| { + Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter.filter) + as Arc<dyn PhysicalExpr> + }) + .collect() + } + + /// Creates a new HashJoinExec with different children while preserving configuration. + /// + /// This method is called during query optimization when the optimizer creates new + /// plan nodes. Importantly, it creates a fresh bounds_accumulator via `try_new` + /// rather than cloning the existing one because partitioning may have changed. + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.builder().with_new_children(children)?.build_exec() + } + ChildrenPropertiesMode::Recompute => self + .builder() + .recompute_properties() + .with_new_children(children)? + .build_exec(), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + self.builder().reset_state().build_exec() + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let on_left = self + .on + .iter() + .map(|on| Arc::clone(&on.0)) + .collect::<Vec<_>>(); + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + + assert_or_internal_err!( + self.mode != PartitionMode::Partitioned + || left_partitions == right_partitions, + "Invalid HashJoinExec, partition count mismatch {left_partitions}!={right_partitions},\ + consider using RepartitionExec" + ); + + assert_or_internal_err!( + self.mode != PartitionMode::CollectLeft || left_partitions == 1, + "Invalid HashJoinExec, the output partition count of the left child must be 1 in CollectLeft mode,\ + consider using CoalescePartitionsExec or the EnforceDistribution rule" + ); + + // Only compute a dynamic filter when the probe subtree contains a consumer. + // Searching from `self` would always find the producer expression owned by this join. + let enable_dynamic_filter_pushdown = if self + .allow_join_dynamic_filter_pushdown(context.session_config().options()) + { + self.dynamic_filter + .as_ref() + .and_then(|df| df.filter.expression_id()) + .map(|id| plan_contains_expression_id(&self.right, id)) + .transpose()? + .unwrap_or(false) + } else { + false + }; + + let join_metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); + + let array_map_created_count = MetricBuilder::new(&self.metrics) + .with_category(MetricCategory::Rows) + .counter(ARRAY_MAP_CREATED_COUNT_METRIC_NAME, partition); + + // Initialize build_accumulator lazily with runtime partition counts (only if enabled) + // Use RepartitionExec's random state (seeds: 0,0,0,0) for partition routing + let repartition_random_state = REPARTITION_RANDOM_STATE; + let build_accumulator = enable_dynamic_filter_pushdown + .then(|| { + self.dynamic_filter.as_ref().map(|df| { + let filter = Arc::clone(&df.filter); + let on_right = self + .on + .iter() + .map(|(_, right_expr)| Arc::clone(right_expr)) + .collect::<Vec<_>>(); + Some(Arc::clone(df.build_accumulator.get_or_init(|| { + Arc::new(SharedBuildAccumulator::new_from_partition_mode( + self.mode, + self.left.as_ref(), + self.right.as_ref(), + filter, + on_right, + repartition_random_state, + self.null_equality, + self.null_aware, + )) + }))) + }) + }) + .flatten() + .flatten(); + + let left_fut = match self.mode { + PartitionMode::CollectLeft => self.left_fut.try_once(|| { + let left_stream = self.left.execute(0, Arc::clone(&context))?; + + let reservation = + MemoryConsumer::new("HashJoinInput").register(context.memory_pool()); + + Ok(collect_left_input( + self.random_state.random_state().clone(), + left_stream, + on_left.clone(), + join_metrics.clone(), + reservation, + need_produce_result_in_final(self.join_type), + self.right().output_partitioning().partition_count(), + enable_dynamic_filter_pushdown, + Arc::clone(context.session_config().options()), + self.null_equality, + array_map_created_count, + )) + })?, + PartitionMode::Partitioned => { + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + + let reservation = + MemoryConsumer::new(format!("HashJoinInput[{partition}]")) + .register(context.memory_pool()); + OnceFut::new(collect_left_input( + self.random_state.random_state().clone(), + left_stream, + on_left.clone(), + join_metrics.clone(), + reservation, + need_produce_result_in_final(self.join_type), + 1, + enable_dynamic_filter_pushdown, + Arc::clone(context.session_config().options()), + self.null_equality, + array_map_created_count, + )) + } + PartitionMode::Auto => { + return plan_err!( + "Invalid HashJoinExec, unsupported PartitionMode {:?} in execute()", + PartitionMode::Auto + ); + } + }; + + let batch_size = context.session_config().batch_size(); + + // we have the batches and the hash map with their keys. We can how create a stream + // over the right that uses this information to issue new batches. + let right_stream = self.right.execute(partition, context)?; + + // update column indices to reflect the projection + let column_indices_after_projection = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|i| self.column_indices[*i].clone()) + .collect(), + None => self.column_indices.clone(), + }; + + let on_right = self + .on + .iter() + .map(|(_, right_expr)| Arc::clone(right_expr)) + .collect::<Vec<_>>(); + + Ok(Box::pin(HashJoinStream::new( + partition, + self.schema(), + on_right, + self.filter.clone(), + self.join_type, + right_stream, + self.random_state.random_state().clone(), + join_metrics, + column_indices_after_projection, + self.null_equality, + HashJoinStreamState::WaitBuildSide, + BuildSide::Initial(BuildSideInitialState { left_fut }), + batch_size, + vec![], + self.right.output_ordering().is_some(), + build_accumulator, + self.mode, + self.null_aware, + self.fetch, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + match (partition, self.mode) { + // Left side is broadcast, so it always needs overall stats + // Right side is partitioned, so it needs per-partition stats + (Some(_), PartitionMode::CollectLeft) => { + vec![ChildStats::At(None), ChildStats::At(partition)] + } + // For Partitioned mode, both sides are hash-partitioned symmetrically, + // so each output partition uses the matching partition from both sides. + (Some(_), PartitionMode::Partitioned) => { + vec![ChildStats::At(partition), ChildStats::At(partition)] + } + // Overall stats requested, look up overall child stats. + (None, _) => vec![ChildStats::At(None), ChildStats::At(None)], + // Auto mode hasn't decided partitioning yet, so it needs + // overall stats from both sides. + (Some(_), PartitionMode::Auto) => { + vec![ChildStats::At(None), ChildStats::At(None)] + } + } + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let left_stats = Arc::clone(&input_stats[0]); + let right_stats = Arc::clone(&input_stats[1]); + let stats = estimate_join_statistics( + Arc::unwrap_or_clone(left_stats), + Arc::unwrap_or_clone(right_stats), + &self.on, + self.null_equality, + &self.join_type, + &self.join_schema, + )?; + // Project statistics if there is a projection + let stats = stats.project(self.projection.as_ref()); + // Apply fetch limit to statistics + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + /// Tries to push `projection` down through `hash_join`. If possible, performs the + /// pushdown and returns a new [`HashJoinExec`] as the top plan which has projections + /// as its children. Otherwise, returns `None`. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // TODO: currently if there is projection in HashJoinExec, we can't push down projection to left or right input. Maybe we can pushdown the mixed projection later. + if self.contains_projection() { + return Ok(None); + } + + let schema = self.schema(); + if let Some(JoinData { + projected_left_child, + projected_right_child, + join_filter, + join_on, + }) = try_pushdown_through_join_with_column_indices( + projection, + self.left(), + self.right(), + self.on(), + &schema, + self.filter(), + self.column_indices.as_slice(), + )? { + self.builder() + .with_new_children(vec![ + Arc::new(projected_left_child), + Arc::new(projected_right_child), + ])? + .with_on(join_on) + .with_filter(join_filter) + // Returned early if projection is not None + .with_projection(None) + .build_exec() + .map(Some) + } else { + try_embed_projection(projection, self) + } + } + + fn gather_filters_for_pushdown( + &self, + phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + config: &ConfigOptions, + ) -> Result<FilterDescription> { + // This is the physical-plan equivalent of `push_down_all_join` in + // `datafusion/optimizer/src/push_down_filter.rs`. That function uses `lr_is_preserved` + // to decide which parent predicates can be pushed past a logical join to its children, + // then checks column references to route each predicate to the correct side. + // + // We apply the same two-level logic here: + // 1. `lr_is_preserved` gates whether a side is eligible at all. + // 2. For each filter, we check that all column references belong to the + // target child (using `column_indices` to map output column positions + // to join sides). This is critical for correctness: name-based matching + // alone (as done by `ChildFilterDescription::from_child`) can incorrectly + // push filters when different join sides have columns with the same name + // (e.g. nested mark joins both producing "mark" columns). + let (left_preserved, right_preserved) = lr_is_preserved(self.join_type); + + // Build the set of allowed column indices for each side + let column_indices: Vec<ColumnIndex> = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|i| self.column_indices[*i].clone()) + .collect(), + None => self.column_indices.clone(), + }; + + let (mut left_allowed, mut right_allowed) = (HashSet::new(), HashSet::new()); + column_indices + .iter() + .enumerate() + .for_each(|(output_idx, ci)| { + match ci.side { + JoinSide::Left => left_allowed.insert(output_idx), + JoinSide::Right => right_allowed.insert(output_idx), + // Mark columns - don't allow pushdown to either side + JoinSide::None => false, + }; + }); + + // For semi joins, filters on output join keys can also be pushed to the + // non-output side: every emitted row has an equal key there. This is not + // true for anti joins, whose emitted rows have no match. + match self.join_type { + JoinType::LeftSemi => { + let left_key_indices: HashSet<usize> = self + .on + .iter() + .filter_map(|(left_key, _)| { + left_key.downcast_ref::<Column>().map(|c| c.index()) + }) + .collect(); + for (output_idx, ci) in column_indices.iter().enumerate() { + if ci.side == JoinSide::Left && left_key_indices.contains(&ci.index) { + right_allowed.insert(output_idx); + } + } + } + JoinType::RightSemi => { + let right_key_indices: HashSet<usize> = self + .on + .iter() + .filter_map(|(_, right_key)| { + right_key.downcast_ref::<Column>().map(|c| c.index()) + }) + .collect(); + for (output_idx, ci) in column_indices.iter().enumerate() { + if ci.side == JoinSide::Right && right_key_indices.contains(&ci.index) + { + left_allowed.insert(output_idx); + } + } + } + _ => {} + } + + let left_child = if left_preserved { + ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + left_allowed, + self.left(), + )? + } else { + ChildFilterDescription::all_unsupported(&parent_filters) + }; + + let mut right_child = if right_preserved { + ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + right_allowed, + self.right(), + )? + } else { + ChildFilterDescription::all_unsupported(&parent_filters) + }; + + // Add dynamic filters in Post phase if enabled. Skip when this join + // already carries a dynamic filter from a previous pass — the shared + // `Arc<DynamicFilterPhysicalExpr>` is still wired into the probe-side + // scan's predicate, and re-creating it would AND a fresh duplicate + // onto every Post-phase invocation (apache/datafusion-ballista#1359 + // surfaces this in AQE replan loops). + if phase == FilterPushdownPhase::Post + && self.dynamic_filter.is_none() + && self.allow_join_dynamic_filter_pushdown(config) + { + // Add actual dynamic filter to right side (probe side) + let dynamic_filter = Self::create_dynamic_filter(&self.on); + right_child = right_child.with_self_filter(dynamic_filter); + } + + Ok(FilterDescription::new() + .with_child(left_child) + .with_child(right_child)) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + let mut result = FilterPushdownPropagation::if_any(child_pushdown_result.clone()); + assert_eq!(child_pushdown_result.self_filters.len(), 2); // Should always be 2, we have 2 children + let right_child_self_filters = &child_pushdown_result.self_filters[1]; // We only push down filters to the right child + // We expect 0 or 1 self filters + if let Some(filter) = right_child_self_filters.first() { + // Note that we don't check PushdDownPredicate::discrimnant because even if nothing said + // "yes, I can fully evaluate this filter" things might still use it for statistics -> it's worth updating + let predicate = Arc::clone(&filter.predicate); + if let Ok(dynamic_filter) = + Arc::downcast::<DynamicFilterPhysicalExpr>(predicate) + { + // We successfully pushed down our self filter - we need to make a new node with the dynamic filter + let new_node = self + .builder() + .with_dynamic_filter(Some(HashJoinExecDynamicFilter { + filter: dynamic_filter, + build_accumulator: OnceLock::new(), + })) + .build_exec()?; + result = result.with_updated_node(new_node); + } + } + Ok(result) + } + + fn supports_limit_pushdown(&self) -> bool { + // Hash join execution plan does not support pushing limit down through to children + // because the children don't know about the join condition and can't + // determine how many rows to produce + false + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + self.builder() + .with_fetch(limit) + .build() + .ok() + .map(|exec| Arc::new(exec) as _) + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let on = self + .on() + .iter() + .map(|(l, r)| -> Result<protobuf::JoinOn> { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(l)?), + right: Some(ctx.encode_expr(r)?), + }) + }) + .collect::<Result<Vec<_>>>()?; + + let join_type = crate::joins::proto::join_type_to_proto(*self.join_type()); + let null_equality = + crate::joins::proto::null_equality_to_proto(self.null_equality()); + // `PartitionMode` is specific to `HashJoinExec`, so its conversion stays + // inline (by-name on purpose: the enums are numbered differently). + let partition_mode = match self.partition_mode() { + PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft, + PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned, + PartitionMode::Auto => protobuf::PartitionMode::Auto, + }; + + let filter = self + .filter() + .map(|f| crate::joins::proto::join_filter_to_proto(f, ctx)) + .transpose()?; + + let dynamic_filter = self + .dynamic_expressions_produced() + .into_iter() + .next() + .map(|expr| ctx.encode_expr(&expr)) + .transpose()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::HashJoin(Box::new( + protobuf::HashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + // Proto3 `repeated` cannot distinguish `None` from + // `Some(vec![])`. `Some(vec![])` (reachable via + // `try_embed_projection` for e.g. `SELECT count(1) … JOIN …`) + // changes the output schema, so it is encoded with the + // single-element sentinel `[u32::MAX]` (never a valid column + // index); every other state is sent as-is. See + // `try_from_proto` for the matching decoder. + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }, + null_aware: self.null_aware, + dynamic_filter, + fetch: self.fetch.map(|f| f as u64), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl HashJoinExec { + /// Reconstruct a [`HashJoinExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_common::{internal_datafusion_err, plan_datafusion_err}; + use datafusion_proto_models::protobuf; + use std::any::Any; + + let hashjoin = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::HashJoin, + "HashJoinExec", + ); + + let left = + ctx.decode_required_child(hashjoin.left.as_deref(), "HashJoinExec", "left")?; + let right = ctx.decode_required_child( + hashjoin.right.as_deref(), + "HashJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = hashjoin + .on + .iter() + .map(|col| { + let l = ctx.decode_required_expr( + col.left.as_ref(), + left_schema.as_ref(), + "HashJoinExec", + "on.left", + )?; + let r = ctx.decode_required_expr( + col.right.as_ref(), + right_schema.as_ref(), + "HashJoinExec", + "on.right", + )?; + Ok((l, r)) + }) + .collect::<Result<_>>()?; + + let join_type = crate::joins::proto::join_type_from_proto( + hashjoin.join_type, + "HashJoinExec", + )?; + let null_equality = crate::joins::proto::null_equality_from_proto( + hashjoin.null_equality, + "HashJoinExec", + )?; + // `PartitionMode` is specific to `HashJoinExec`, so its conversion stays + // inline (by-name on purpose: the enums are numbered differently). + let partition_mode = match protobuf::PartitionMode::try_from( + hashjoin.partition_mode, + ) + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec: unknown PartitionMode {}", + hashjoin.partition_mode + ) + })? { + protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft, + protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned, + protobuf::PartitionMode::Auto => PartitionMode::Auto, + }; + + let filter = hashjoin + .filter + .as_ref() + .map(|f| crate::joins::proto::join_filter_from_proto(f, ctx, "HashJoinExec")) + .transpose()?; + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match hashjoin.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + // Restore the row limit that `limit_pushdown` may have pushed into the + // join. The field is presence-tracked, so a message written before it + // existed decodes to `None` (no limit) rather than to `Some(0)`. + // + // The conversion is checked, not `as usize`: `fetch` is a `u64` on the + // wire but a `usize` in the plan, and on a 32-bit target `as usize` + // truncates. A fetch of `1 << 32` would become `0` -- not merely a + // wrong limit but the worst one, silently turning the query into an + // empty result. Report the out-of-range value instead. Please do not + // "simplify" this back to `as usize`. + let fetch = hashjoin + .fetch + .map(|f| { + usize::try_from(f).map_err(|_| { + plan_datafusion_err!( + "HashJoinExec: fetch value {f} cannot be represented as usize on this target" + ) + }) + }) + .transpose()?; + + let mut hash_join = HashJoinExecBuilder::new(left, right, on, join_type) + .with_filter(filter) + .with_projection(projection) + .with_partition_mode(partition_mode) + .with_null_equality(null_equality) + .with_null_aware(hashjoin.null_aware) + .with_fetch(fetch) + .build()?; + + if let Some(dynamic_filter_proto) = &hashjoin.dynamic_filter { + // The dynamic filter is a `DynamicFilterPhysicalExpr` over the probe + // (right) side; decode against the right schema then downcast. + let dynamic_filter_expr = + ctx.decode_expr(dynamic_filter_proto, right_schema.as_ref())?; + let df = (dynamic_filter_expr as Arc<dyn Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + hash_join = hash_join.with_dynamic_filter_expr(df)?; + } + + Ok(Arc::new(hash_join)) + } +} + +/// Determines which sides of a join are "preserved" for filter pushdown. +/// +/// A preserved side means filters on that side's columns can be safely pushed +/// below the join. This mostly mirrors the logical optimizer's `lr_is_preserved`; +/// semi joins additionally allow join-key filters on the non-output side. +fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { + match join_type { + JoinType::Inner => (true, true), + JoinType::Left => (true, false), + JoinType::Right => (false, true), + JoinType::Full => (false, false), + // Callers restrict the non-output side of semi joins to join-key columns. + JoinType::LeftSemi | JoinType::RightSemi => (true, true), + JoinType::LeftAnti | JoinType::LeftMark => (true, false), + JoinType::RightAnti | JoinType::RightMark => (false, true), + } +} + +/// Accumulator for collecting min/max bounds from build-side data during hash join. +/// +/// This struct encapsulates the logic for progressively computing column bounds +/// (minimum and maximum values) for a specific join key expression as batches +/// are processed during the build phase of a hash join. +/// +/// The bounds are used for dynamic filter pushdown optimization, where filters +/// based on the actual data ranges can be pushed down to the probe side to +/// eliminate unnecessary data early. +struct CollectLeftAccumulator { + /// The physical expression to evaluate for each batch + expr: Arc<dyn PhysicalExpr>, + /// Accumulator for tracking the minimum value across all batches + min: MinAccumulator, + /// Accumulator for tracking the maximum value across all batches + max: MaxAccumulator, +} + +impl CollectLeftAccumulator { + /// Creates a new accumulator for tracking bounds of a join key expression. + /// + /// # Arguments + /// * `expr` - The physical expression to track bounds for + /// * `schema` - The schema of the input data + /// + /// # Returns + /// A new `CollectLeftAccumulator` instance configured for the expression's data type + fn try_new(expr: Arc<dyn PhysicalExpr>, schema: &SchemaRef) -> Result<Self> { + /// Recursively unwraps dictionary types to get the underlying value type. + fn dictionary_value_type(data_type: &DataType) -> DataType { + match data_type { + DataType::Dictionary(_, value_type) => { + dictionary_value_type(value_type.as_ref()) + } + _ => data_type.clone(), + } + } + + let data_type = expr + .data_type(schema) + // Min/Max can operate on dictionary data but expect to be initialized with the underlying value type + .map(|dt| dictionary_value_type(&dt))?; + Ok(Self { + expr, + min: MinAccumulator::try_new(&data_type)?, + max: MaxAccumulator::try_new(&data_type)?, + }) + } + + /// Updates the accumulators with values from a new batch. + /// + /// Evaluates the expression on the batch and updates both min and max + /// accumulators with the resulting values. + /// + /// # Arguments + /// * `batch` - The record batch to process + /// + /// # Returns + /// Ok(()) if the update succeeds, or an error if expression evaluation fails + fn update_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let array = self.expr.evaluate(batch)?.into_array(batch.num_rows())?; + self.min.update_batch(std::slice::from_ref(&array))?; + self.max.update_batch(std::slice::from_ref(&array))?; + Ok(()) + } + + /// Finalizes the accumulation and returns the computed bounds. + /// + /// Consumes self to extract the final min and max values from the accumulators. + /// + /// # Returns + /// The `ColumnBounds` containing the minimum and maximum values observed + fn evaluate(mut self) -> Result<ColumnBounds> { + Ok(ColumnBounds::new( + self.min.evaluate()?, + self.max.evaluate()?, + )) + } +} + +/// State for collecting the build-side data during hash join +struct BuildSideState { + batches: Vec<RecordBatch>, + num_rows: usize, + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + bounds_accumulators: Option<Vec<CollectLeftAccumulator>>, + /// Counts the memory of `batches` for `reservation`. Batches can share + /// underlying buffers (e.g. when the input emits zero-copy slices of one + /// larger batch), so each buffer must be reserved only once. + memory_counter: RecordBatchMemoryCounter, +} + +impl BuildSideState { + /// Create a new BuildSideState with optional accumulators for bounds computation + fn try_new( + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + on_left: Vec<Arc<dyn PhysicalExpr>>, + schema: &SchemaRef, + should_compute_dynamic_filters: bool, + ) -> Result<Self> { + Ok(Self { + batches: Vec::new(), + num_rows: 0, + metrics, + reservation, + memory_counter: RecordBatchMemoryCounter::new(), + bounds_accumulators: should_compute_dynamic_filters + .then(|| { + on_left + .into_iter() + .map(|expr| CollectLeftAccumulator::try_new(expr, schema)) + .collect::<Result<Vec<_>>>() + }) + .transpose()?, + }) + } +} + +fn should_collect_min_max_for_perfect_hash( + on_left: &[PhysicalExprRef], + schema: &SchemaRef, +) -> Result<bool> { + if on_left.len() != 1 { + return Ok(false); + } + + let expr = &on_left[0]; + let data_type = expr.data_type(schema)?; + Ok(ArrayMap::is_supported_type(&data_type)) +} + +/// Collects all batches from the left (build) side stream and creates a hash map for joining. +/// +/// This function is responsible for: +/// 1. Consuming the entire left stream and collecting all batches into memory +/// 2. Building a hash map from the join key columns for efficient probe operations +/// 3. Computing bounds for dynamic filter pushdown (if enabled) +/// 4. Preparing visited indices bitmap for certain join types +/// +/// # Parameters +/// * `random_state` - Random state for consistent hashing across partitions +/// * `left_stream` - Stream of record batches from the build side +/// * `on_left` - Physical expressions for the left side join keys +/// * `metrics` - Metrics collector for tracking memory usage and row counts +/// * `reservation` - Memory reservation tracker for the hash table and data +/// * `with_visited_indices_bitmap` - Whether to track visited indices (for outer joins) +/// * `probe_threads_count` - Number of threads that will probe this hash table +/// * `should_compute_dynamic_filters` - Whether to compute min/max bounds for dynamic filtering +/// +/// # Dynamic Filter Coordination +/// When `should_compute_dynamic_filters` is true, this function computes the min/max bounds +/// for each join key column but does NOT update the dynamic filter. Instead, the +/// bounds are stored in the returned `JoinLeftData` and later coordinated by +/// `SharedBuildAccumulator` to ensure all partitions contribute their bounds +/// before updating the filter exactly once. +/// +/// # Returns +/// `JoinLeftData` containing the hash map, consolidated batch, join key values, +/// visited indices bitmap, and computed bounds (if requested). +#[expect(clippy::too_many_arguments)] +async fn collect_left_input( + random_state: RandomState, + left_stream: SendableRecordBatchStream, + on_left: Vec<PhysicalExprRef>, + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + with_visited_indices_bitmap: bool, + probe_threads_count: usize, + should_compute_dynamic_filters: bool, + config: Arc<ConfigOptions>, + null_equality: NullEquality, + array_map_created_count: Count, +) -> Result<JoinLeftData> { + let schema = left_stream.schema(); + + let should_collect_min_max_for_phj = + should_collect_min_max_for_perfect_hash(&on_left, &schema)?; + + let initial = BuildSideState::try_new( + metrics, + reservation, + on_left.clone(), + &schema, + should_compute_dynamic_filters || should_collect_min_max_for_phj, + )?; + + let state = left_stream + .try_fold(initial, |mut state, batch| async move { + // Update accumulators if computing bounds + if let Some(ref mut accumulators) = state.bounds_accumulators { + for accumulator in accumulators { + accumulator.update_batch(&batch)?; + } + } + + // Decide if we spill or not + let batch_size = state.memory_counter.count_batch(&batch); + // Reserve memory for incoming batch + state.reservation.try_grow(batch_size)?; + // Update metrics + state.metrics.build_mem_used.add(batch_size); + state.metrics.build_input_batches.add(1); + state.metrics.build_input_rows.add(batch.num_rows()); + // Update row count + state.num_rows += batch.num_rows(); + // Push batch to output + state.batches.push(batch); + Ok(state) + }) + .await?; + + // Extract fields from state + let BuildSideState { + batches, + num_rows, + metrics, + mut reservation, + bounds_accumulators, + memory_counter: _, + } = state; + + // Compute bounds + let mut bounds = match bounds_accumulators { + Some(accumulators) if num_rows > 0 => { + let bounds = accumulators + .into_iter() + .map(CollectLeftAccumulator::evaluate) + .collect::<Result<Vec<_>>>()?; + Some(PartitionBounds::new(bounds)) + } + _ => None, + }; + + let (join_hash_map, batch, left_values) = + if let Some((array_map, batch, left_value)) = try_create_array_map( + &bounds, + &schema, + &batches, + &on_left, + &mut reservation, + config.execution.perfect_hash_join_small_build_threshold, + config.execution.perfect_hash_join_min_key_density, + null_equality, + )? { + array_map_created_count.add(1); + metrics.build_mem_used.add(array_map.size()); + + (Map::ArrayMap(array_map), batch, left_value) + } else { + // Estimation of memory size, required for hashtable, prior to allocation. + // Final result can be verified using `RawTable.allocation_info()` + let fixed_size_u32 = size_of::<JoinHashMapU32>(); + let fixed_size_u64 = size_of::<JoinHashMapU64>(); + + // Use `u32` indices for the JoinHashMap when num_rows ≤ u32::MAX, otherwise use the + // `u64` indice variant + // Arc is used instead of Box to allow sharing with SharedBuildAccumulator for hash map pushdown + let mut hashmap: Box<dyn JoinHashMapType> = if num_rows > u32::MAX as usize { + let estimated_hashtable_size = + estimate_memory_size::<(u64, u64)>(num_rows, fixed_size_u64)?; + reservation.try_grow(estimated_hashtable_size)?; + metrics.build_mem_used.add(estimated_hashtable_size); + Box::new(JoinHashMapU64::with_capacity(num_rows)) + } else { + let estimated_hashtable_size = + estimate_memory_size::<(u32, u64)>(num_rows, fixed_size_u32)?; + reservation.try_grow(estimated_hashtable_size)?; + metrics.build_mem_used.add(estimated_hashtable_size); + Box::new(JoinHashMapU32::with_capacity(num_rows)) + }; + + let mut hashes_buffer = Vec::new(); + let mut offset = 0; + + let batches_iter = batches.iter().rev(); + + // Updating hashmap starting from the last batch + for batch in batches_iter.clone() { + hashes_buffer.clear(); + hashes_buffer.resize(batch.num_rows(), 0); + update_hash( + &on_left, + batch, + &mut *hashmap, + offset, + &random_state, + &mut hashes_buffer, + 0, + true, + null_equality, + )?; + offset += batch.num_rows(); + } + + // Merge all batches into a single batch, so we can directly index into the arrays + let batch = concat_batches(&schema, batches_iter.clone())?; + + let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?; + + (Map::HashMap(hashmap), batch, left_values) + }; + + // Reserve additional memory for visited indices bitmap and create shared builder + let visited_indices_bitmap = if with_visited_indices_bitmap { + let bitmap_size = bit_util::ceil(batch.num_rows(), 8); + reservation.try_grow(bitmap_size)?; + metrics.build_mem_used.add(bitmap_size); + + let mut bitmap_buffer = BooleanBufferBuilder::new(batch.num_rows()); + bitmap_buffer.append_n(num_rows, false); + bitmap_buffer + } else { + BooleanBufferBuilder::new(0) + }; + + let map = Arc::new(join_hash_map); + + // Nothing reads the strategy unless this join computes dynamic filters. + let membership = if num_rows == 0 || !should_compute_dynamic_filters { + PushdownStrategy::Empty + } else { + // If the build side is small enough we can use IN list pushdown. + // If it's too big we fall back to pushing down a reference to the hash table. + // See `PushdownStrategy` for more details. + let estimated_size = left_values + .iter() + .map(|arr| arr.get_array_memory_size()) + .sum::<usize>(); + + let pushdown_inlist = !left_values.is_empty() + && !left_values[0].is_empty() + && estimated_size <= config.optimizer.hash_join_inlist_pushdown_max_size + && map.num_of_distinct_key() + <= config + .optimizer + .hash_join_inlist_pushdown_max_distinct_values; + + if pushdown_inlist + && let Some(in_list_values) = build_struct_inlist_values(&left_values)? + { + PushdownStrategy::InList(in_list_values) + } else { + // Past the InList threshold use a bucket bitmap for container pruning. + let pruning_bitmap = match (left_values.as_slice(), bounds.as_ref()) { + ([keys], Some(bounds)) if !keys.is_empty() => bounds + .get_column_bounds(0) + .and_then(|b| { + KeyRangeBitmap::try_new( + keys, + &b.min, + &b.max, + map.num_of_distinct_key(), + ) + }) + .map(Arc::new), + _ => None, + }; + // Held for the join's lifetime, so charge it like the maps; it is + // optional, so skip it rather than fail when the pool is full. + let pruning_bitmap = pruning_bitmap.filter(|bitmap| { + let ok = reservation.try_grow(bitmap.size()).is_ok(); + if ok { + metrics.build_mem_used.add(bitmap.size()); + } + ok + }); + PushdownStrategy::Map(Arc::clone(&map), pruning_bitmap) + } + }; + + if should_collect_min_max_for_phj && !should_compute_dynamic_filters { + bounds = None; + } + + let data = JoinLeftData { + map, + batch, + values: left_values, + visited_indices_bitmap: Mutex::new(visited_indices_bitmap), + probe_threads_counter: AtomicUsize::new(probe_threads_count), + _reservation: reservation, + bounds, + membership, + probe_side_non_empty: AtomicBool::new(false), + probe_side_has_null: AtomicBool::new(false), + }; + + Ok(data) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_phj_used(metrics: &MetricsSet, use_phj: bool) { + if use_phj { + assert!( + metrics + .sum_by_name(ARRAY_MAP_CREATED_COUNT_METRIC_NAME) + .expect("should have array_map_created_count metrics") + .as_usize() + >= 1 + ); + } else { + assert_eq!( + metrics + .sum_by_name(ARRAY_MAP_CREATED_COUNT_METRIC_NAME) + .map(|v| v.as_usize()) + .unwrap_or(0), + 0 + ) + } + } + + fn build_schema_and_on() -> Result<(SchemaRef, SchemaRef, JoinOn)> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("b1", DataType::Int32, true), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, true), + Field::new("b1", DataType::Int32, true), + ])); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left_schema)?) as _, + Arc::new(Column::new_with_schema("b1", &right_schema)?) as _, + )]; + Ok((left_schema, right_schema, on)) + } + + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::execution_plan::Boundedness; + use crate::filter::FilterExecBuilder; + use crate::joins::hash_join::stream::lookup_join_hashmap; + use crate::test::{TestMemoryExec, assert_join_metrics}; + use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; + use crate::{ + common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, + test::exec::MockExec, + }; + + use arrow::array::{ + Date32Array, Int32Array, Int64Array, StructArray, UInt32Array, UInt64Array, + }; + use arrow::buffer::NullBuffer; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::hash_utils::create_hashes; + use datafusion_common::test_util::{batches_to_sort_string, batches_to_string}; + use datafusion_common::{ + ScalarValue, assert_batches_eq, assert_batches_sorted_eq, assert_contains, + exec_err, internal_err, + }; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion_physical_expr::{ + EquivalenceProperties, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; + use hashbrown::HashTable; + use insta::{allow_duplicates, assert_snapshot}; + use rstest::*; + use rstest_reuse::*; + + /// Runs a join whose 200 keys spread over 2M would size a pruning bitmap + /// at the 128 KiB cap, and reports the bytes the build side reserved. + /// + /// Forces the IN-list threshold off: under `force_hash_collisions` every + /// key hashes to the same bucket, so `num_of_distinct_key()` no longer + /// reflects the real 200 distinct keys and would otherwise take the + /// IN-list branch instead of the one under test. + async fn build_mem_used(limit: usize, dynamic_filters: bool) -> Result<usize> { + let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int64, false)])); + let batch = |keys: Vec<i64>| { + let column = Arc::new(Int64Array::from(keys)) as ArrayRef; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![column])?; + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) + }; + let on = vec![( + Arc::new(Column::new_with_schema("k", &schema)?) as _, + Arc::new(Column::new_with_schema("k", &schema)?) as _, + )]; + let build = batch((0..200).map(|i| i * 10_000).collect())?; + let probe = batch(vec![0, 500_000, 1_500_000])?; + let join = if dynamic_filters { + hash_join_with_dynamic_filter(build, probe, on, JoinType::Inner)?.0 + } else { + join( + build, + probe, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )? + }; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(limit, 1.0) + .build_arc()?; + let mut config = SessionConfig::new(); + config + .options_mut() + .optimizer + .hash_join_inlist_pushdown_max_size = 0; + let task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(config), + ); + let batches = common::collect(join.execute(0, task_ctx)?).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 3); + Ok(join + .metrics() + .unwrap() + .sum_by_name("build_mem_used") + .unwrap() + .as_usize()) + } + + /// The bitmap is pruning-only: never built when no dynamic filter will read + /// it, and dropped rather than fatal when the pool cannot fit it. + #[tokio::test] + async fn pruning_bitmap_is_optional() -> Result<()> { + assert!(build_mem_used(1_000_000, false).await? < 100_000); + assert!(build_mem_used(1_000_000, true).await? > 100_000); + build_mem_used(100_000, true).await?; + Ok(()) + } + + #[derive(Debug)] + struct PartitionedTestExec { + cache: Arc<PlanProperties>, + } + + impl PartitionedTestExec { + fn try_new(schema: SchemaRef, partitioning: Partitioning) -> Result<Self> { + Ok(Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )), + }) + } + } + + impl DisplayAs for PartitionedTestExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PartitionedTestExec") + } + } + + impl ExecutionPlan for PartitionedTestExec { + fn name(&self) -> &'static str { + "PartitionedTestExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unreachable!() + } + } + + fn div_ceil(a: usize, b: usize) -> usize { + a.div_ceil(b) + } + + #[template] + #[rstest] + fn hash_join_exec_configs( + #[values(8192, 10, 5, 2, 1)] batch_size: usize, + #[values(true, false)] use_perfect_hash_join_as_possible: bool, + ) { + } + + fn prepare_task_ctx( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Arc<TaskContext> { + let mut session_config = SessionConfig::default().with_batch_size(batch_size); + + if use_perfect_hash_join_as_possible { + session_config + .options_mut() + .execution + .perfect_hash_join_small_build_threshold = 819200; + session_config + .options_mut() + .execution + .perfect_hash_join_min_key_density = 0.0; + } else { + session_config + .options_mut() + .execution + .perfect_hash_join_small_build_threshold = 0; + session_config + .options_mut() + .execution + .perfect_hash_join_min_key_density = f64::INFINITY; + } + Arc::new(TaskContext::default().with_session_config(session_config)) + } + + fn build_table( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), + ) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + /// Build a table with two columns supporting nullable values + fn build_table_two_cols( + a: (&str, &Vec<Option<i32>>), + b: (&str, &Vec<Option<i32>>), + ) -> Arc<dyn ExecutionPlan> { + let schema = Arc::new(Schema::new(vec![ + Field::new(a.0, DataType::Int32, true), + Field::new(b.0, DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(a.1.clone())), + Arc::new(Int32Array::from(b.1.clone())), + ], + ) + .unwrap(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + fn join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: &JoinType, + null_equality: NullEquality, + ) -> Result<HashJoinExec> { + HashJoinExec::try_new( + left, + right, + on, + None, + join_type, + None, + PartitionMode::CollectLeft, + null_equality, + false, + ) + } + + fn join_with_filter( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: JoinFilter, + join_type: &JoinType, + null_equality: NullEquality, + ) -> Result<HashJoinExec> { + HashJoinExec::try_new( + left, + right, + on, + Some(filter), + join_type, + None, + PartitionMode::CollectLeft, + null_equality, + false, + ) + } + + fn empty_build_with_probe_error_inputs() + -> (Arc<dyn ExecutionPlan>, Arc<dyn ExecutionPlan>, JoinOn) { + let left_batch = + build_table_i32(("a1", &vec![]), ("b1", &vec![]), ("c1", &vec![])); + let left_schema = left_batch.schema(); + let left: Arc<dyn ExecutionPlan> = TestMemoryExec::try_new_exec( + &[vec![left_batch]], + Arc::clone(&left_schema), + None, + ) + .unwrap(); + + let err = exec_err!("bad data error"); + let right_batch = + build_table_i32(("a2", &vec![]), ("b1", &vec![]), ("c2", &vec![])); + let right_schema = right_batch.schema(); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left_schema).unwrap()) as _, + Arc::new(Column::new_with_schema("b1", &right_schema).unwrap()) as _, + )]; + let right: Arc<dyn ExecutionPlan> = Arc::new( + MockExec::new(vec![Ok(right_batch), err], right_schema) + .with_use_task(false) + // The planted error must only surface if the probe side is + // polled, not when a parent node computes statistics during + // planning. + .with_unknown_statistics(), + ); + + (left, right, on) + } + + async fn assert_empty_build_probe_behavior( + join_types: &[JoinType], + expect_probe_error: bool, + with_filter: bool, + ) { + let (left, right, on) = empty_build_with_probe_error_inputs(); + let filter = prepare_join_filter(); + + for join_type in join_types { + let join = if with_filter { + join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + join_type, + NullEquality::NullEqualsNothing, + ) + .unwrap() + } else { + join( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + NullEquality::NullEqualsNothing, + ) + .unwrap() + }; + + let result = common::collect( + join.execute(0, Arc::new(TaskContext::default())).unwrap(), + ) + .await; + + if expect_probe_error { + let result_string = result.unwrap_err().to_string(); + assert!( + result_string.contains("bad data error"), + "actual: {result_string}" + ); + } else { + let batches = result.unwrap(); + assert!( + batches.is_empty(), + "expected no output batches for {join_type}, got {batches:?}" + ); + } + } + } + + fn hash_join_with_dynamic_filter( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, + ) -> Result<(HashJoinExec, Arc<DynamicFilterPhysicalExpr>)> { + hash_join_with_dynamic_filter_and_mode( + left, + right, + on, + join_type, + PartitionMode::CollectLeft, + ) + } + + fn hash_join_with_dynamic_filter_and_mode( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, + mode: PartitionMode, + ) -> Result<(HashJoinExec, Arc<DynamicFilterPhysicalExpr>)> { + let dynamic_filter = HashJoinExec::create_dynamic_filter(&on); + let consumer: Arc<dyn PhysicalExpr> = Arc::clone(&dynamic_filter) as _; + let right = Arc::new(FilterExecBuilder::new(consumer, right).build()?); + let mut join = HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + mode, + NullEquality::NullEqualsNothing, + false, + )?; + join.dynamic_filter = Some(HashJoinExecDynamicFilter { + filter: Arc::clone(&dynamic_filter), + build_accumulator: OnceLock::new(), + }); + + Ok((join, dynamic_filter)) + } + + async fn join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: &JoinType, + null_equality: NullEquality, + context: Arc<TaskContext>, + ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> { + let join = join(left, right, on, join_type, null_equality)?; + let columns_header = columns(&join.schema()); + + let stream = join.execute(0, context)?; + let batches = common::collect(stream).await?; + let metrics = join.metrics().unwrap(); + + Ok((columns_header, batches, metrics)) + } + + async fn partitioned_join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: &JoinType, + null_equality: NullEquality, + context: Arc<TaskContext>, + ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> { + join_collect_with_partition_mode( + left, + right, + on, + join_type, + PartitionMode::Partitioned, + null_equality, + context, + ) + .await + } + + async fn join_collect_with_partition_mode( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: &JoinType, + partition_mode: PartitionMode, + null_equality: NullEquality, + context: Arc<TaskContext>, + ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> { + let partition_count = 4; + + let (left_expr, right_expr) = on + .iter() + .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) + .unzip(); + + let left_repartitioned: Arc<dyn ExecutionPlan> = match partition_mode { + PartitionMode::CollectLeft => Arc::new(CoalescePartitionsExec::new(left)), + PartitionMode::Partitioned => Arc::new(RepartitionExec::try_new( + left, + Partitioning::Hash(left_expr, partition_count), + )?), + PartitionMode::Auto => { + return internal_err!("Unexpected PartitionMode::Auto in join tests"); + } + }; + + let right_repartitioned: Arc<dyn ExecutionPlan> = match partition_mode { + PartitionMode::CollectLeft => { + let partition_column_name = right.schema().field(0).name().clone(); + let partition_expr = vec![Arc::new(Column::new_with_schema( + &partition_column_name, + &right.schema(), + )?) as _]; + Arc::new(RepartitionExec::try_new( + right, + Partitioning::Hash(partition_expr, partition_count), + )?) as _ + } + PartitionMode::Partitioned => Arc::new(RepartitionExec::try_new( + right, + Partitioning::Hash(right_expr, partition_count), + )?), + PartitionMode::Auto => { + return internal_err!("Unexpected PartitionMode::Auto in join tests"); + } + }; + + let join = HashJoinExec::try_new( + left_repartitioned, + right_repartitioned, + on, + None, + join_type, + None, + partition_mode, + null_equality, + false, + )?; + + let columns = columns(&join.schema()); + + let mut batches = vec![]; + for i in 0..partition_count { + let stream = join.execute(i, Arc::clone(&context))?; + let more_batches = common::collect(stream).await?; + batches.extend( + more_batches + .into_iter() + .filter(|b| b.num_rows() > 0) + .collect::<Vec<_>>(), + ); + } + let metrics = join.metrics().unwrap(); + + Ok((columns, batches, metrics)) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_inner_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 5]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + allow_duplicates! { + // Inner join output is expected to preserve both inputs order + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 5 | 9 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn partitioned_join_inner_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 5]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = partitioned_join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 5 | 9 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[tokio::test] + async fn join_inner_one_no_shared_column_names() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 5]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b2", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 5 | 9 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[tokio::test] + async fn join_inner_one_randomly_ordered() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = build_table( + ("a1", &vec![0, 3, 2, 1]), + ("b1", &vec![4, 5, 5, 4]), + ("c1", &vec![6, 9, 8, 7]), + ); + let right = build_table( + ("a2", &vec![20, 30, 10]), + ("b2", &vec![5, 6, 4]), + ("c2", &vec![80, 90, 70]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 3 | 5 | 9 | 20 | 5 | 80 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 0 | 4 | 6 | 10 | 4 | 70 | + | 1 | 4 | 7 | 10 | 4 | 70 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 4); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_inner_two( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b2", &vec![1, 2, 2]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a1", &vec![1, 2, 3]), + ("b2", &vec![1, 2, 2]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b2", "c1", "a1", "b2", "c2"]); + + let expected_batch_count = if cfg!(not(feature = "force_hash_collisions")) { + // Expected number of hash table matches = 3 + // in case batch_size is 1 - additional empty batch for remaining 3-2 row + let mut expected_batch_count = div_ceil(3, batch_size); + if batch_size == 1 { + expected_batch_count += 1; + } + expected_batch_count + } else { + // With hash collisions enabled, all records will match each other + // and filtered later. + div_ceil(9, batch_size) + }; + + // With batch coalescing, we may have fewer batches than expected + assert!( + batches.len() <= expected_batch_count, + "expected at most {expected_batch_count} batches, got {}", + batches.len() + ); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 1 | 7 | 1 | 1 | 70 | + | 2 | 2 | 8 | 2 | 2 | 80 | + | 2 | 2 | 9 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + /// Test where the left has 2 parts, the right with 1 part => 1 part + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_inner_one_two_parts_left( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let batch1 = build_table_i32( + ("a1", &vec![1, 2]), + ("b2", &vec![1, 2]), + ("c1", &vec![7, 8]), + ); + let batch2 = + build_table_i32(("a1", &vec![2]), ("b2", &vec![2]), ("c1", &vec![9])); + let schema = batch1.schema(); + let left = + TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) + .unwrap(); + let left = Arc::new(CoalescePartitionsExec::new(left)); + + let right = build_table( + ("a1", &vec![1, 2, 3]), + ("b2", &vec![1, 2, 2]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b2", "c1", "a1", "b2", "c2"]); + + let expected_batch_count = if cfg!(not(feature = "force_hash_collisions")) { + // Expected number of hash table matches = 3 + // in case batch_size is 1 - additional empty batch for remaining 3-2 row + let mut expected_batch_count = div_ceil(3, batch_size); + if batch_size == 1 { + expected_batch_count += 1; + } + expected_batch_count + } else { + // With hash collisions enabled, all records will match each other + // and filtered later. + div_ceil(9, batch_size) + }; + + // With batch coalescing, we may have fewer batches than expected + assert!( + batches.len() <= expected_batch_count, + "expected at most {expected_batch_count} batches, got {}", + batches.len() + ); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 1 | 7 | 1 | 1 | 70 | + | 2 | 2 | 8 | 2 | 2 | 80 | + | 2 | 2 | 9 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[tokio::test] + async fn join_inner_one_two_parts_left_randomly_ordered() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch1 = build_table_i32( + ("a1", &vec![0, 3]), + ("b1", &vec![4, 5]), + ("c1", &vec![6, 9]), + ); + let batch2 = build_table_i32( + ("a1", &vec![2, 1]), + ("b1", &vec![5, 4]), + ("c1", &vec![8, 7]), + ); + let schema = batch1.schema(); + + let left = + TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) + .unwrap(); + let left = Arc::new(CoalescePartitionsExec::new(left)); + let right = build_table( + ("a2", &vec![20, 30, 10]), + ("b2", &vec![5, 6, 4]), + ("c2", &vec![80, 90, 70]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 3 | 5 | 9 | 20 | 5 | 80 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 0 | 4 | 6 | 10 | 4 | 70 | + | 1 | 4 | 7 | 10 | 4 | 70 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 4); + + Ok(()) + } + + /// Test where the left has 1 part, the right has 2 parts => 2 parts + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_inner_one_two_parts_right( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 5]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + + let batch1 = build_table_i32( + ("a2", &vec![10, 20]), + ("b1", &vec![4, 6]), + ("c2", &vec![70, 80]), + ); + let batch2 = + build_table_i32(("a2", &vec![30]), ("b1", &vec![5]), ("c2", &vec![90])); + let schema = batch1.schema(); + let right = + TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) + .unwrap(); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + // first part + let stream = join.execute(0, Arc::clone(&task_ctx))?; + let batches = common::collect(stream).await?; + + let expected_batch_count = if cfg!(not(feature = "force_hash_collisions")) { + // Expected number of hash table matches for first right batch = 1 + // and additional empty batch for non-joined 20-6-80 + let mut expected_batch_count = div_ceil(1, batch_size); + if batch_size == 1 { + expected_batch_count += 1; + } + expected_batch_count + } else { + // With hash collisions enabled, all records will match each other + // and filtered later. + div_ceil(6, batch_size) + }; + // With batch coalescing, we may have fewer batches than expected + assert!( + batches.len() <= expected_batch_count, + "expected at most {expected_batch_count} batches, got {}", + batches.len() + ); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + +----+----+----+----+----+----+ + "); + } + + // second part + let stream = join.execute(1, Arc::clone(&task_ctx))?; + let batches = common::collect(stream).await?; + + let expected_batch_count = if cfg!(not(feature = "force_hash_collisions")) { + // Expected number of hash table matches for second right batch = 2 + div_ceil(2, batch_size) + } else { + // With hash collisions enabled, all records will match each other + // and filtered later. + div_ceil(3, batch_size) + }; + // With batch coalescing, we may have fewer batches than expected + assert!( + batches.len() <= expected_batch_count, + "expected at most {expected_batch_count} batches, got {}", + batches.len() + ); + + // Inner join output is expected to preserve both inputs order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 2 | 5 | 8 | 30 | 5 | 90 | + | 3 | 5 | 9 | 30 | 5 | 90 | + +----+----+----+----+----+----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + fn build_table_two_batches( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), + ) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch.clone(), batch]], schema, None).unwrap() + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_multi_batch( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table_two_batches( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b1", &right.schema()).unwrap()) as _, + )]; + + let join = join( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Left, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + let (_, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Left, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + return Ok(()); + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_full_multi_batch( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + // create two identical batches for the right side + let right = build_table_two_batches( + ("a2", &vec![10, 20, 30]), + ("b2", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::Full, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + let stream = join.execute(0, task_ctx).unwrap(); + let batches = common::collect(stream).await.unwrap(); + let metrics = join.metrics().unwrap(); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 30 | 6 | 90 | + | | | | 30 | 6 | 90 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_empty_right( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table_i32(("a2", &vec![]), ("b1", &vec![]), ("c2", &vec![])); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b1", &right.schema()).unwrap()) as _, + )]; + let schema = right.schema(); + let right = TestMemoryExec::try_new_exec(&[vec![right]], schema, None).unwrap(); + let join = join( + left, + right, + on, + &JoinType::Left, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + let stream = join.execute(0, task_ctx).unwrap(); + let batches = common::collect(stream).await.unwrap(); + let metrics = join.metrics().unwrap(); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | | | | + | 2 | 5 | 8 | | | | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_full_empty_right( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table_i32(("a2", &vec![]), ("b2", &vec![]), ("c2", &vec![])); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + let schema = right.schema(); + let right = TestMemoryExec::try_new_exec(&[vec![right]], schema, None).unwrap(); + let join = join( + left, + right, + on, + &JoinType::Full, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + let stream = join.execute(0, task_ctx).unwrap(); + let batches = common::collect(stream).await.unwrap(); + let metrics = join.metrics().unwrap(); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | | | | + | 2 | 5 | 8 | | | | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Left, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + /// Under NullEqualsNothing, NULL join keys are not inserted into the hash + /// map, so a build side whose keys are all NULL produces an empty map even + /// though it contains rows. Join types that emit unmatched build rows must + /// still produce them from the visited bitmap. + #[rstest] + #[tokio::test] + async fn join_all_null_build_keys( + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + partition_mode: PartitionMode, + ) -> Result<()> { + let left = build_table_two_cols( + ("a1", &vec![Some(1), Some(2)]), + ("b1", &vec![None, None]), // all build-side join keys are NULL + ); + let right = build_table_two_cols( + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), None, Some(6)]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + let (_, batches, metrics) = join_collect_with_partition_mode( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &join_type, + partition_mode, + NullEquality::NullEqualsNothing, + Arc::new(TaskContext::default()), + ) + .await?; + + // For join types whose output requires a build-side match, an + // empty map guarantees an empty result, so `state_after_build_ready` + // completes the stream without ever fetching a probe batch (probe + // `input_rows` stays 0). All other join types must still scan the + // probe side. `input_rows` is summed across every partition. + let probe_rows = metrics + .sum_by_name("input_rows") + .map(|v| v.as_usize()) + .unwrap_or(0); + if join_type.empty_map_produces_empty_result() { + assert_eq!( + probe_rows, 0, + "{join_type} should skip the probe side for an all-NULL build" + ); + } else { + assert!(probe_rows > 0, "{join_type} must scan the probe side"); + } + + match join_type { + JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi => { + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, 0, "unexpected rows for {join_type}"); + } + JoinType::Left => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | 1 | | | | + | 2 | | | | + +----+----+----+----+ + "); + } + } + JoinType::Right => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | | | 10 | 4 | + | | | 20 | | + | | | 30 | 6 | + +----+----+----+----+ + "); + } + } + JoinType::Full => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | | | 10 | 4 | + | | | 20 | | + | | | 30 | 6 | + | 1 | | | | + | 2 | | | | + +----+----+----+----+ + "); + } + } + JoinType::LeftAnti => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+ + | a1 | b1 | + +----+----+ + | 1 | | + | 2 | | + +----+----+ + "); + } + } + JoinType::RightAnti => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+ + | a2 | b1 | + +----+----+ + | 10 | 4 | + | 20 | | + | 30 | 6 | + +----+----+ + "); + } + } + JoinType::LeftMark => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-------+ + | a1 | b1 | mark | + +----+----+-------+ + | 1 | | false | + | 2 | | false | + +----+----+-------+ + "); + } + } + JoinType::RightMark => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-------+ + | a2 | b1 | mark | + +----+----+-------+ + | 10 | 4 | false | + | 20 | | false | + | 30 | 6 | false | + +----+----+-------+ + "); + } + } + } + } + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn partitioned_join_left_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = partitioned_join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Left, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + fn build_semi_anti_left_table() -> Arc<dyn ExecutionPlan> { + // just two line match + // b1 = 10 + build_table( + ("a1", &vec![1, 3, 5, 7, 9, 11, 13]), + ("b1", &vec![1, 3, 5, 7, 8, 8, 10]), + ("c1", &vec![10, 30, 50, 70, 90, 110, 130]), + ) + } + + fn build_semi_anti_right_table() -> Arc<dyn ExecutionPlan> { + // just two line match + // b2 = 10 + build_table( + ("a2", &vec![8, 12, 6, 2, 10, 4]), + ("b2", &vec![8, 10, 6, 2, 10, 4]), + ("c2", &vec![20, 40, 60, 80, 100, 120]), + ) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_semi( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + // left_table left semi join right_table on left_table.b1 = right_table.b2 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::LeftSemi, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // ignore the order + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a1 | b1 | c1 | + +----+----+-----+ + | 11 | 8 | 110 | + | 13 | 10 | 130 | + | 9 | 8 | 90 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_semi_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + + // left_table left semi join right_table on left_table.b1 = right_table.b2 and right_table.a2 != 10 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let column_indices = vec![ColumnIndex { + index: 0, + side: JoinSide::Right, + }]; + let intermediate_schema = + Schema::new(vec![Field::new("x", DataType::Int32, true)]); + + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices.clone(), + Arc::new(intermediate_schema.clone()), + ); + + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter, + &JoinType::LeftSemi, + NullEquality::NullEqualsNothing, + )?; + + let columns_header = columns(&join.schema()); + assert_eq!(columns_header.clone(), vec!["a1", "b1", "c1"]); + + let stream = join.execute(0, Arc::clone(&task_ctx))?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a1 | b1 | c1 | + +----+----+-----+ + | 11 | 8 | 110 | + | 13 | 10 | 130 | + | 9 | 8 | 90 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + // left_table left semi join right_table on left_table.b1 = right_table.b2 and right_table.a2 > 10 + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + let filter = JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema), + ); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::LeftSemi, + NullEquality::NullEqualsNothing, + )?; + + let columns_header = columns(&join.schema()); + assert_eq!(columns_header, vec!["a1", "b1", "c1"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a1 | b1 | c1 | + +----+----+-----+ + | 13 | 10 | 130 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_semi( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + + // left_table right semi join right_table on left_table.b1 = right_table.b2 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::RightSemi, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a2", "b2", "c2"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // RightSemi join output is expected to preserve right input order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 8 | 8 | 20 | + | 12 | 10 | 40 | + | 10 | 10 | 100 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_semi_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + + // left_table right semi join right_table on left_table.b1 = right_table.b2 on left_table.a1!=9 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let column_indices = vec![ColumnIndex { + index: 0, + side: JoinSide::Left, + }]; + let intermediate_schema = + Schema::new(vec![Field::new("x", DataType::Int32, true)]); + + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(9)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices.clone(), + Arc::new(intermediate_schema.clone()), + ); + + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter, + &JoinType::RightSemi, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a2", "b2", "c2"]); + + let stream = join.execute(0, Arc::clone(&task_ctx))?; + let batches = common::collect(stream).await?; + + // RightSemi join output is expected to preserve right input order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 8 | 8 | 20 | + | 12 | 10 | 40 | + | 10 | 10 | 100 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + // left_table right semi join right_table on left_table.b1 = right_table.b2 on left_table.a1!=9 + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(11)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema.clone()), + ); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::RightSemi, + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // RightSemi join output is expected to preserve right input order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 12 | 10 | 40 | + | 10 | 10 | 100 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_anti( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + // left_table left anti join right_table on left_table.b1 = right_table.b2 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::LeftAnti, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 1 | 10 | + | 3 | 3 | 30 | + | 5 | 5 | 50 | + | 7 | 7 | 70 | + +----+----+----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_anti_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + // left_table left anti join right_table on left_table.b1 = right_table.b2 and right_table.a2!=8 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let column_indices = vec![ColumnIndex { + index: 0, + side: JoinSide::Right, + }]; + let intermediate_schema = + Schema::new(vec![Field::new("x", DataType::Int32, true)]); + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(8)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices.clone(), + Arc::new(intermediate_schema.clone()), + ); + + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter, + &JoinType::LeftAnti, + NullEquality::NullEqualsNothing, + )?; + + let columns_header = columns(&join.schema()); + assert_eq!(columns_header, vec!["a1", "b1", "c1"]); + + let stream = join.execute(0, Arc::clone(&task_ctx))?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a1 | b1 | c1 | + +----+----+-----+ + | 1 | 1 | 10 | + | 11 | 8 | 110 | + | 3 | 3 | 30 | + | 5 | 5 | 50 | + | 7 | 7 | 70 | + | 9 | 8 | 90 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + // left_table left anti join right_table on left_table.b1 = right_table.b2 and right_table.a2 != 13 + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(8)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema), + ); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::LeftAnti, + NullEquality::NullEqualsNothing, + )?; + + let columns_header = columns(&join.schema()); + assert_eq!(columns_header, vec!["a1", "b1", "c1"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a1 | b1 | c1 | + +----+----+-----+ + | 1 | 1 | 10 | + | 11 | 8 | 110 | + | 3 | 3 | 30 | + | 5 | 5 | 50 | + | 7 | 7 | 70 | + | 9 | 8 | 90 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_anti( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::RightAnti, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a2", "b2", "c2"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // RightAnti join output is expected to preserve right input order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 6 | 6 | 60 | + | 2 | 2 | 80 | + | 4 | 4 | 120 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_anti_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_semi_anti_left_table(); + let right = build_semi_anti_right_table(); + // left_table right anti join right_table on left_table.b1 = right_table.b2 and left_table.a1!=13 + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let column_indices = vec![ColumnIndex { + index: 0, + side: JoinSide::Left, + }]; + let intermediate_schema = + Schema::new(vec![Field::new("x", DataType::Int32, true)]); + + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(13)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema.clone()), + ); + + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter, + &JoinType::RightAnti, + NullEquality::NullEqualsNothing, + )?; + + let columns_header = columns(&join.schema()); + assert_eq!(columns_header, vec!["a2", "b2", "c2"]); + + let stream = join.execute(0, Arc::clone(&task_ctx))?; + let batches = common::collect(stream).await?; + + // RightAnti join output is expected to preserve right input order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 12 | 10 | 40 | + | 6 | 6 | 60 | + | 2 | 2 | 80 | + | 10 | 10 | 100 | + | 4 | 4 | 120 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + // left_table right anti join right_table on left_table.b1 = right_table.b2 and right_table.b2!=8 + let column_indices = vec![ColumnIndex { + index: 1, + side: JoinSide::Right, + }]; + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(8)))), + )) as Arc<dyn PhysicalExpr>; + + let filter = JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema), + ); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::RightAnti, + NullEquality::NullEqualsNothing, + )?; + + let columns_header = columns(&join.schema()); + assert_eq!(columns_header, vec!["a2", "b2", "c2"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // RightAnti join output is expected to preserve right input order + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 8 | 8 | 20 | + | 6 | 6 | 60 | + | 2 | 2 | 80 | + | 4 | 4 | 120 | + +----+----+-----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Right, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | | | | 30 | 6 | 90 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn partitioned_join_right_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = partitioned_join_collect( + left, + right, + on, + &JoinType::Right, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b1", "c2"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | | | | 30 | 6 | 90 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_full_one( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b2", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::Full, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 30 | 6 | 90 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_mark( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::LeftMark, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "mark"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+-------+ + | a1 | b1 | c1 | mark | + +----+----+----+-------+ + | 1 | 4 | 7 | true | + | 2 | 5 | 8 | true | + | 3 | 7 | 9 | false | + +----+----+----+-------+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn partitioned_join_left_mark( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40]), + ("b1", &vec![4, 4, 5, 6]), + ("c2", &vec![60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = partitioned_join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::LeftMark, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "mark"]); + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+-------+ + | a1 | b1 | c1 | mark | + +----+----+----+-------+ + | 1 | 4 | 7 | true | + | 2 | 5 | 8 | true | + | 3 | 7 | 9 | false | + +----+----+----+-------+ + "); + } + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_mark( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::RightMark, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a2", "b1", "c2", "mark"]); + + let expected = [ + "+----+----+----+-------+", + "| a2 | b1 | c2 | mark |", + "+----+----+----+-------+", + "| 10 | 4 | 70 | true |", + "| 20 | 5 | 80 | true |", + "| 30 | 6 | 90 | false |", + "+----+----+----+-------+", + ]; + assert_batches_sorted_eq!(expected, &batches); + + assert_join_metrics!(metrics, 3); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn partitioned_join_right_mark( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40]), + ("b1", &vec![4, 4, 5, 6]), // 6 does not exist on the left + ("c2", &vec![60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = partitioned_join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::RightMark, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a2", "b1", "c2", "mark"]); + + let expected = [ + "+----+----+----+-------+", + "| a2 | b1 | c2 | mark |", + "+----+----+----+-------+", + "| 10 | 4 | 60 | true |", + "| 20 | 4 | 70 | true |", + "| 30 | 5 | 80 | true |", + "| 40 | 6 | 90 | false |", + "+----+----+----+-------+", + ]; + assert_batches_sorted_eq!(expected, &batches); + + assert_join_metrics!(metrics, 4); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[test] + fn join_with_hash_collisions_64() -> Result<()> { + let mut hashmap_left = HashTable::with_capacity(4); + let left = build_table_i32( + ("a", &vec![10, 20]), + ("x", &vec![100, 200]), + ("y", &vec![200, 300]), + ); + + let random_state = RandomState::with_seed(0); + let hashes_buff = &mut vec![0; left.num_rows()]; + let hashes = create_hashes([&left.columns()[0]], &random_state, hashes_buff)?; + + // Maps both values to both indices (1 and 2, representing input 0 and 1) + // 0 -> (0, 1) + // 1 -> (0, 2) + // The equality check will make sure only hashes[0] maps to 0 and hashes[1] maps to 1 + hashmap_left.insert_unique(hashes[0], (hashes[0], 1), |(h, _)| *h); + hashmap_left.insert_unique(hashes[0], (hashes[0], 2), |(h, _)| *h); + + hashmap_left.insert_unique(hashes[1], (hashes[1], 1), |(h, _)| *h); + hashmap_left.insert_unique(hashes[1], (hashes[1], 2), |(h, _)| *h); + + let next = vec![2, 0]; + + let right = build_table_i32( + ("a", &vec![10, 20]), + ("b", &vec![0, 0]), + ("c", &vec![30, 40]), + ); + + // Join key column for both join sides + let key_column: PhysicalExprRef = Arc::new(Column::new("a", 0)) as _; + + let join_hash_map = JoinHashMapU64::new(hashmap_left, next); + + let left_keys_values = key_column.evaluate(&left)?.into_array(left.num_rows())?; + let right_keys_values = + key_column.evaluate(&right)?.into_array(right.num_rows())?; + let mut hashes_buffer = vec![0; right.num_rows()]; + create_hashes([&right_keys_values], &random_state, &mut hashes_buffer)?; + + let mut probe_indices_buffer = Vec::new(); + let mut build_indices_buffer = Vec::new(); + let (l, r, _) = lookup_join_hashmap( + &join_hash_map, + &[left_keys_values], + &[right_keys_values], + NullEquality::NullEqualsNothing, + &hashes_buffer, + None, + 8192, + (0, None), + &mut probe_indices_buffer, + &mut build_indices_buffer, + )?; + + let left_ids: UInt64Array = vec![0, 1].into(); + + let right_ids: UInt32Array = vec![0, 1].into(); + + assert_eq!(left_ids, l); + + assert_eq!(right_ids, r); + + Ok(()) + } + + #[test] + fn join_with_hash_collisions_u32() -> Result<()> { + let mut hashmap_left = HashTable::with_capacity(4); + let left = build_table_i32( + ("a", &vec![10, 20]), + ("x", &vec![100, 200]), + ("y", &vec![200, 300]), + ); + + let random_state = RandomState::with_seed(0); + let hashes_buff = &mut vec![0; left.num_rows()]; + let hashes = create_hashes([&left.columns()[0]], &random_state, hashes_buff)?; + + hashmap_left.insert_unique(hashes[0], (hashes[0], 1u32), |(h, _)| *h); + hashmap_left.insert_unique(hashes[0], (hashes[0], 2u32), |(h, _)| *h); + hashmap_left.insert_unique(hashes[1], (hashes[1], 1u32), |(h, _)| *h); + hashmap_left.insert_unique(hashes[1], (hashes[1], 2u32), |(h, _)| *h); + + let next: Vec<u32> = vec![2, 0]; + + let right = build_table_i32( + ("a", &vec![10, 20]), + ("b", &vec![0, 0]), + ("c", &vec![30, 40]), + ); + + let key_column: PhysicalExprRef = Arc::new(Column::new("a", 0)) as _; + + let join_hash_map = JoinHashMapU32::new(hashmap_left, next); + + let left_keys_values = key_column.evaluate(&left)?.into_array(left.num_rows())?; + let right_keys_values = + key_column.evaluate(&right)?.into_array(right.num_rows())?; + let mut hashes_buffer = vec![0; right.num_rows()]; + create_hashes([&right_keys_values], &random_state, &mut hashes_buffer)?; + + let mut probe_indices_buffer = Vec::new(); + let mut build_indices_buffer = Vec::new(); + let (l, r, _) = lookup_join_hashmap( + &join_hash_map, + &[left_keys_values], + &[right_keys_values], + NullEquality::NullEqualsNothing, + &hashes_buffer, + None, + 8192, + (0, None), + &mut probe_indices_buffer, + &mut build_indices_buffer, + )?; + + // We still expect to match rows 0 and 1 on both sides + let left_ids: UInt64Array = vec![0, 1].into(); + let right_ids: UInt32Array = vec![0, 1].into(); + + assert_eq!(left_ids, l); + assert_eq!(right_ids, r); + + Ok(()) + } + + #[tokio::test] + async fn join_with_duplicated_column_names() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = build_table( + ("a", &vec![1, 2, 3]), + ("b", &vec![4, 5, 7]), + ("c", &vec![7, 8, 9]), + ); + let right = build_table( + ("a", &vec![10, 20, 30]), + ("b", &vec![1, 2, 7]), + ("c", &vec![70, 80, 90]), + ); + let on = vec![( + // join on a=b so there are duplicate column names on unjoined columns + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a", "b", "c", "a", "b", "c"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +---+---+---+----+---+----+ + | a | b | c | a | b | c | + +---+---+---+----+---+----+ + | 1 | 4 | 7 | 10 | 1 | 70 | + | 2 | 5 | 8 | 20 | 2 | 80 | + +---+---+---+----+---+----+ + "); + } + + Ok(()) + } + + fn prepare_join_filter() -> JoinFilter { + let column_indices = vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ]; + let intermediate_schema = Schema::new(vec![ + Field::new("c", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ]); + let filter_expression = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c", 0)), + Operator::Gt, + Arc::new(Column::new("c", 1)), + )) as Arc<dyn PhysicalExpr>; + + JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema), + ) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_inner_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a", &vec![0, 1, 2, 2]), + ("b", &vec![4, 5, 7, 8]), + ("c", &vec![7, 8, 9, 1]), + ); + let right = build_table( + ("a", &vec![10, 20, 30, 40]), + ("b", &vec![2, 2, 3, 4]), + ("c", &vec![7, 5, 6, 4]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + let filter = prepare_join_filter(); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a", "b", "c", "a", "b", "c"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +---+---+---+----+---+---+ + | a | b | c | a | b | c | + +---+---+---+----+---+---+ + | 2 | 7 | 9 | 10 | 2 | 7 | + | 2 | 7 | 9 | 20 | 2 | 5 | + +---+---+---+----+---+---+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_left_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a", &vec![0, 1, 2, 2]), + ("b", &vec![4, 5, 7, 8]), + ("c", &vec![7, 8, 9, 1]), + ); + let right = build_table( + ("a", &vec![10, 20, 30, 40]), + ("b", &vec![2, 2, 3, 4]), + ("c", &vec![7, 5, 6, 4]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + let filter = prepare_join_filter(); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::Left, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a", "b", "c", "a", "b", "c"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +---+---+---+----+---+---+ + | a | b | c | a | b | c | + +---+---+---+----+---+---+ + | 0 | 4 | 7 | | | | + | 1 | 5 | 8 | | | | + | 2 | 7 | 9 | 10 | 2 | 7 | + | 2 | 7 | 9 | 20 | 2 | 5 | + | 2 | 8 | 1 | | | | + +---+---+---+----+---+---+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_right_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a", &vec![0, 1, 2, 2]), + ("b", &vec![4, 5, 7, 8]), + ("c", &vec![7, 8, 9, 1]), + ); + let right = build_table( + ("a", &vec![10, 20, 30, 40]), + ("b", &vec![2, 2, 3, 4]), + ("c", &vec![7, 5, 6, 4]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + let filter = prepare_join_filter(); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::Right, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a", "b", "c", "a", "b", "c"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +---+---+---+----+---+---+ + | a | b | c | a | b | c | + +---+---+---+----+---+---+ + | | | | 30 | 3 | 6 | + | | | | 40 | 4 | 4 | + | 2 | 7 | 9 | 10 | 2 | 7 | + | 2 | 7 | 9 | 20 | 2 | 5 | + +---+---+---+----+---+---+ + "); + } + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn join_full_with_filter( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let left = build_table( + ("a", &vec![0, 1, 2, 2]), + ("b", &vec![4, 5, 7, 8]), + ("c", &vec![7, 8, 9, 1]), + ); + let right = build_table( + ("a", &vec![10, 20, 30, 40]), + ("b", &vec![2, 2, 3, 4]), + ("c", &vec![7, 5, 6, 4]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b", &right.schema()).unwrap()) as _, + )]; + let filter = prepare_join_filter(); + + let join = join_with_filter( + left, + right, + on, + filter, + &JoinType::Full, + NullEquality::NullEqualsNothing, + )?; + + let columns = columns(&join.schema()); + assert_eq!(columns, vec!["a", "b", "c", "a", "b", "c"]); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let expected = [ + "+---+---+---+----+---+---+", + "| a | b | c | a | b | c |", + "+---+---+---+----+---+---+", + "| | | | 30 | 3 | 6 |", + "| | | | 40 | 4 | 4 |", + "| 2 | 7 | 9 | 10 | 2 | 7 |", + "| 2 | 7 | 9 | 20 | 2 | 5 |", + "| 0 | 4 | 7 | | | |", + "| 1 | 5 | 8 | | | |", + "| 2 | 8 | 1 | | | |", + "+---+---+---+----+---+---+", + ]; + assert_batches_sorted_eq!(expected, &batches); + + let metrics = join.metrics().unwrap(); + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + // THIS MIGRATION HALTED DUE TO ISSUE #15312 + //allow_duplicates! { + // assert_snapshot!(batches_to_sort_string(&batches), @r#" + // +---+---+---+----+---+---+ + // | a | b | c | a | b | c | + // +---+---+---+----+---+---+ + // | | | | 30 | 3 | 6 | + // | | | | 40 | 4 | 4 | + // | 2 | 7 | 9 | 10 | 2 | 7 | + // | 2 | 7 | 9 | 20 | 2 | 5 | + // | 0 | 4 | 7 | | | | + // | 1 | 5 | 8 | | | | + // | 2 | 8 | 1 | | | | + // +---+---+---+----+---+---+ + // "#) + //} + + Ok(()) + } + + /// Test for parallelized HashJoinExec with PartitionMode::CollectLeft + #[tokio::test] + async fn test_collect_left_multiple_partitions_join() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b2", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + + let expected_inner = vec![ + "+----+----+----+----+----+----+", + "| a1 | b1 | c1 | a2 | b2 | c2 |", + "+----+----+----+----+----+----+", + "| 1 | 4 | 7 | 10 | 4 | 70 |", + "| 2 | 5 | 8 | 20 | 5 | 80 |", + "+----+----+----+----+----+----+", + ]; + let expected_left = vec![ + "+----+----+----+----+----+----+", + "| a1 | b1 | c1 | a2 | b2 | c2 |", + "+----+----+----+----+----+----+", + "| 1 | 4 | 7 | 10 | 4 | 70 |", + "| 2 | 5 | 8 | 20 | 5 | 80 |", + "| 3 | 7 | 9 | | | |", + "+----+----+----+----+----+----+", + ]; + let expected_right = vec![ + "+----+----+----+----+----+----+", + "| a1 | b1 | c1 | a2 | b2 | c2 |", + "+----+----+----+----+----+----+", + "| | | | 30 | 6 | 90 |", + "| 1 | 4 | 7 | 10 | 4 | 70 |", + "| 2 | 5 | 8 | 20 | 5 | 80 |", + "+----+----+----+----+----+----+", + ]; + let expected_full = vec![ + "+----+----+----+----+----+----+", + "| a1 | b1 | c1 | a2 | b2 | c2 |", + "+----+----+----+----+----+----+", + "| | | | 30 | 6 | 90 |", + "| 1 | 4 | 7 | 10 | 4 | 70 |", + "| 2 | 5 | 8 | 20 | 5 | 80 |", + "| 3 | 7 | 9 | | | |", + "+----+----+----+----+----+----+", + ]; + let expected_left_semi = vec![ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "| 1 | 4 | 7 |", + "| 2 | 5 | 8 |", + "+----+----+----+", + ]; + let expected_left_anti = vec![ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "| 3 | 7 | 9 |", + "+----+----+----+", + ]; + let expected_right_semi = vec![ + "+----+----+----+", + "| a2 | b2 | c2 |", + "+----+----+----+", + "| 10 | 4 | 70 |", + "| 20 | 5 | 80 |", + "+----+----+----+", + ]; + let expected_right_anti = vec![ + "+----+----+----+", + "| a2 | b2 | c2 |", + "+----+----+----+", + "| 30 | 6 | 90 |", + "+----+----+----+", + ]; + let expected_left_mark = vec![ + "+----+----+----+-------+", + "| a1 | b1 | c1 | mark |", + "+----+----+----+-------+", + "| 1 | 4 | 7 | true |", + "| 2 | 5 | 8 | true |", + "| 3 | 7 | 9 | false |", + "+----+----+----+-------+", + ]; + let expected_right_mark = vec![ + "+----+----+----+-------+", + "| a2 | b2 | c2 | mark |", + "+----+----+----+-------+", + "| 10 | 4 | 70 | true |", + "| 20 | 5 | 80 | true |", + "| 30 | 6 | 90 | false |", + "+----+----+----+-------+", + ]; + + let test_cases = vec![ + (JoinType::Inner, expected_inner), + (JoinType::Left, expected_left), + (JoinType::Right, expected_right), + (JoinType::Full, expected_full), + (JoinType::LeftSemi, expected_left_semi), + (JoinType::LeftAnti, expected_left_anti), + (JoinType::RightSemi, expected_right_semi), + (JoinType::RightAnti, expected_right_anti), + (JoinType::LeftMark, expected_left_mark), + (JoinType::RightMark, expected_right_mark), + ]; + + for (join_type, expected) in test_cases { + let (_, batches, metrics) = join_collect_with_partition_mode( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &join_type, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + Arc::clone(&task_ctx), + ) + .await?; + assert_batches_sorted_eq!(expected, &batches); + assert_join_metrics!(metrics, expected.len() - 4); + } + + Ok(()) + } + + #[tokio::test] + async fn join_date32() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("date", DataType::Date32, false), + Field::new("n", DataType::Int32, false), + ])); + + let dates: ArrayRef = Arc::new(Date32Array::from(vec![19107, 19108, 19109])); + let n: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![dates, n])?; + let left = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) + .unwrap(); + let dates: ArrayRef = Arc::new(Date32Array::from(vec![19108, 19108, 19109])); + let n: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![dates, n])?; + let right = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("date", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("date", &right.schema()).unwrap()) as _, + )]; + + let join = join( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )?; + + let task_ctx = Arc::new(TaskContext::default()); + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +------------+---+------------+---+ + | date | n | date | n | + +------------+---+------------+---+ + | 2022-04-26 | 2 | 2022-04-26 | 4 | + | 2022-04-26 | 2 | 2022-04-26 | 5 | + | 2022-04-27 | 3 | 2022-04-27 | 6 | + +------------+---+------------+---+ + "); + } + + Ok(()) + } + + #[tokio::test] + async fn join_with_error_right() { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + + // right input stream returns one good batch and then one error. + // The error should be returned. + let err = exec_err!("bad data error"); + let right = build_table_i32(("a2", &vec![]), ("b1", &vec![]), ("c2", &vec![])); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b1", &right.schema()).unwrap()) as _, + )]; + let schema = right.schema(); + let right = build_table_i32(("a2", &vec![]), ("b1", &vec![]), ("c2", &vec![])); + let right_input = Arc::new(MockExec::new(vec![Ok(right), err], schema)); + + let join_types = vec![ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + ]; + + for join_type in join_types { + let join = join( + Arc::clone(&left), + Arc::clone(&right_input) as Arc<dyn ExecutionPlan>, + on.clone(), + &join_type, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + let task_ctx = Arc::new(TaskContext::default()); + + let stream = join.execute(0, task_ctx).unwrap(); + + // Expect that an error is returned + let result_string = common::collect(stream).await.unwrap_err().to_string(); + assert!( + result_string.contains("bad data error"), + "actual: {result_string}" + ); + } + } + + #[tokio::test] + async fn join_does_not_consume_probe_when_empty_build_fixes_output() { + assert_empty_build_probe_behavior( + &[ + JoinType::Inner, + JoinType::Left, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightSemi, + ], + false, + false, + ) + .await; + } + + #[tokio::test] + async fn join_does_not_consume_probe_when_empty_build_fixes_output_with_filter() { + assert_empty_build_probe_behavior( + &[ + JoinType::Inner, + JoinType::Left, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightSemi, + ], + false, + true, + ) + .await; + } + + #[tokio::test] + async fn join_still_consumes_probe_when_empty_build_needs_probe_rows() { + assert_empty_build_probe_behavior( + &[ + JoinType::Right, + JoinType::Full, + JoinType::RightAnti, + JoinType::RightMark, + ], + true, + false, + ) + .await; + } + + #[tokio::test] + async fn join_still_consumes_probe_when_empty_build_needs_probe_rows_with_filter() { + assert_empty_build_probe_behavior( + &[ + JoinType::Right, + JoinType::Full, + JoinType::RightAnti, + JoinType::RightMark, + ], + true, + true, + ) + .await; + } + + #[tokio::test] + async fn join_split_batch() { + let left = build_table( + ("a1", &vec![1, 2, 3, 4]), + ("b1", &vec![1, 1, 1, 1]), + ("c1", &vec![0, 0, 0, 0]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40, 50]), + ("b2", &vec![1, 1, 1, 1, 1]), + ("c2", &vec![0, 0, 0, 0, 0]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + + let join_types = vec![ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::LeftSemi, + JoinType::LeftAnti, + ]; + let expected_resultset_records = 20; + let common_result = [ + "+----+----+----+----+----+----+", + "| a1 | b1 | c1 | a2 | b2 | c2 |", + "+----+----+----+----+----+----+", + "| 1 | 1 | 0 | 10 | 1 | 0 |", + "| 2 | 1 | 0 | 10 | 1 | 0 |", + "| 3 | 1 | 0 | 10 | 1 | 0 |", + "| 4 | 1 | 0 | 10 | 1 | 0 |", + "| 1 | 1 | 0 | 20 | 1 | 0 |", + "| 2 | 1 | 0 | 20 | 1 | 0 |", + "| 3 | 1 | 0 | 20 | 1 | 0 |", + "| 4 | 1 | 0 | 20 | 1 | 0 |", + "| 1 | 1 | 0 | 30 | 1 | 0 |", + "| 2 | 1 | 0 | 30 | 1 | 0 |", + "| 3 | 1 | 0 | 30 | 1 | 0 |", + "| 4 | 1 | 0 | 30 | 1 | 0 |", + "| 1 | 1 | 0 | 40 | 1 | 0 |", + "| 2 | 1 | 0 | 40 | 1 | 0 |", + "| 3 | 1 | 0 | 40 | 1 | 0 |", + "| 4 | 1 | 0 | 40 | 1 | 0 |", + "| 1 | 1 | 0 | 50 | 1 | 0 |", + "| 2 | 1 | 0 | 50 | 1 | 0 |", + "| 3 | 1 | 0 | 50 | 1 | 0 |", + "| 4 | 1 | 0 | 50 | 1 | 0 |", + "+----+----+----+----+----+----+", + ]; + let left_batch = [ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "| 1 | 1 | 0 |", + "| 2 | 1 | 0 |", + "| 3 | 1 | 0 |", + "| 4 | 1 | 0 |", + "+----+----+----+", + ]; + let right_batch = [ + "+----+----+----+", + "| a2 | b2 | c2 |", + "+----+----+----+", + "| 10 | 1 | 0 |", + "| 20 | 1 | 0 |", + "| 30 | 1 | 0 |", + "| 40 | 1 | 0 |", + "| 50 | 1 | 0 |", + "+----+----+----+", + ]; + let right_empty = [ + "+----+----+----+", + "| a2 | b2 | c2 |", + "+----+----+----+", + "+----+----+----+", + ]; + let left_empty = [ + "+----+----+----+", + "| a1 | b1 | c1 |", + "+----+----+----+", + "+----+----+----+", + ]; + + // validation of partial join results output for different batch_size setting + for join_type in join_types { + for batch_size in (1..21).rev() { + let task_ctx = prepare_task_ctx(batch_size, true); + + let join = join( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &join_type, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + let stream = join.execute(0, task_ctx).unwrap(); + let batches = common::collect(stream).await.unwrap(); + + // For inner/right join expected batch count equals dev_ceil result, + // as there is no need to append non-joined build side data. + // For other join types it'll be div_ceil + 1 -- for additional batch + // containing not visited build side rows (empty in this test case). + let expected_batch_count = match join_type { + JoinType::Inner + | JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti => { + div_ceil(expected_resultset_records, batch_size) + } + _ => div_ceil(expected_resultset_records, batch_size) + 1, + }; + // With batch coalescing, we may have fewer batches than expected + assert!( + batches.len() <= expected_batch_count, + "expected at most {expected_batch_count} output batches for {join_type} join with batch_size = {batch_size}, got {}", + batches.len() + ); + + let expected = match join_type { + JoinType::RightSemi => right_batch.to_vec(), + JoinType::RightAnti => right_empty.to_vec(), + JoinType::LeftSemi => left_batch.to_vec(), + JoinType::LeftAnti => left_empty.to_vec(), + _ => common_result.to_vec(), + }; + // For anti joins with empty results, we may get zero batches + // (with coalescing) instead of one empty batch with schema + if batches.is_empty() { + // Verify this is an expected empty result case + assert!( + matches!(join_type, JoinType::RightAnti | JoinType::LeftAnti), + "Unexpected empty result for {join_type} join" + ); + } else { + assert_batches_eq!(expected, &batches); + } + } + } + } + + #[tokio::test] + async fn single_partition_join_overallocation() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("b1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("c1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ); + let right = build_table( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + + let join_types = vec![ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ]; + + for join_type in join_types { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + let join = join( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &join_type, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let err = common::collect(stream).await.unwrap_err(); + + // Asserting that operator-level reservation attempting to overallocate + assert_contains!( + err.to_string(), + "Resources exhausted: Additional allocation failed for HashJoinInput with top memory consumers (across reservations) as:\n HashJoinInput" + ); + + assert_contains!( + err.to_string(), + "Failed to allocate additional 120.0 B for HashJoinInput" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn build_side_sliced_batches_memory_accounting() -> Result<()> { + // The build side emits zero-copy slices of one large batch, as e.g. an + // aggregate emitting its output in batch_size chunks does. The buffers + // shared by the slices must be reserved once in total, not once per + // slice: per-slice accounting reserves number_of_slices x parent size + // and aborts queries that fit in memory with room to spare. + let n = 4096; + let v: Vec<i32> = (0..n).collect(); + let parent = build_table_i32(("a1", &v), ("b1", &v), ("c1", &v)); + let slices: Vec<RecordBatch> = + (0..16).map(|i| parent.slice(i * 256, 256)).collect(); + let left = + TestMemoryExec::try_new_exec(&[slices], parent.schema(), None).unwrap(); + + let right_batch = build_table_i32( + ("a2", &vec![10, 11]), + ("b2", &vec![0, 1]), + ("c2", &vec![14, 15]), + ); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch.clone()]], + right_batch.schema(), + None, + ) + .unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &parent.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right_batch.schema())?) as _, + )]; + + // Enough for the parent batch (~48KB) plus the join hash table, but far + // below the ~768KB that per-slice accounting would reserve + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(400_000, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + let join = join( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, 2); + + Ok(()) + } + + #[tokio::test] + async fn partitioned_join_overallocation() -> Result<()> { + // Prepare partitioned inputs for HashJoinExec + // No need to adjust partitioning, as execution should fail with `Resources exhausted` error + let left_batch = build_table_i32( + ("a1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("b1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("c1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ); + let left = TestMemoryExec::try_new_exec( + &[vec![left_batch.clone()], vec![left_batch.clone()]], + left_batch.schema(), + None, + ) + .unwrap(); + let right_batch = build_table_i32( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + ); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch.clone()], vec![right_batch.clone()]], + right_batch.schema(), + None, + ) + .unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left_batch.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right_batch.schema())?) as _, + )]; + + let join_types = vec![ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + ]; + + for join_type in join_types { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .build_arc()?; + let session_config = SessionConfig::default().with_batch_size(50); + let task_ctx = TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + let join = HashJoinExec::try_new( + Arc::clone(&left) as Arc<dyn ExecutionPlan>, + Arc::clone(&right) as Arc<dyn ExecutionPlan>, + on.clone(), + None, + &join_type, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + + let stream = join.execute(1, task_ctx)?; + let err = common::collect(stream).await.unwrap_err(); + + // Asserting that stream-level reservation attempting to overallocate + assert_contains!( + err.to_string(), + "Resources exhausted: Additional allocation failed for HashJoinInput[1] with top memory consumers (across reservations) as:\n HashJoinInput[1]" + ); + + assert_contains!( + err.to_string(), + "Failed to allocate additional 120.0 B for HashJoinInput[1]" + ); + } + + Ok(()) + } + + fn build_table_struct( + struct_name: &str, + field_name_and_values: (&str, &Vec<Option<i32>>), + nulls: Option<NullBuffer>, + ) -> Arc<dyn ExecutionPlan> { + let (field_name, values) = field_name_and_values; + let inner_fields = vec![Field::new(field_name, DataType::Int32, true)]; + let schema = Schema::new(vec![Field::new( + struct_name, + DataType::Struct(inner_fields.clone().into()), + nulls.is_some(), + )]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(StructArray::new( + inner_fields.into(), + vec![Arc::new(Int32Array::from(values.clone()))], + nulls, + ))], + ) + .unwrap(); + let schema_ref = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema_ref, None).unwrap() + } + + #[tokio::test] + async fn join_on_struct() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = + build_table_struct("n1", ("a", &vec![None, Some(1), Some(2), Some(3)]), None); + let right = + build_table_struct("n2", ("a", &vec![None, Some(1), Some(2), Some(4)]), None); + let on = vec![( + Arc::new(Column::new_with_schema("n1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("n2", &right.schema())?) as _, + )]; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["n1", "n2"]); + + allow_duplicates! { + assert_snapshot!(batches_to_string(&batches), @r" + +--------+--------+ + | n1 | n2 | + +--------+--------+ + | {a: } | {a: } | + | {a: 1} | {a: 1} | + | {a: 2} | {a: 2} | + +--------+--------+ + "); + } + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[tokio::test] + async fn join_on_struct_with_nulls() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = + build_table_struct("n1", ("a", &vec![None]), Some(NullBuffer::new_null(1))); + let right = + build_table_struct("n2", ("a", &vec![None]), Some(NullBuffer::new_null(1))); + let on = vec![( + Arc::new(Column::new_with_schema("n1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("n2", &right.schema())?) as _, + )]; + + let (_, batches_null_eq, metrics) = join_collect( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &JoinType::Inner, + NullEquality::NullEqualsNull, + Arc::clone(&task_ctx), + ) + .await?; + + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches_null_eq), @r" + +----+----+ + | n1 | n2 | + +----+----+ + | | | + +----+----+ + "); + } + + assert_join_metrics!(metrics, 1); + + let (_, batches_null_neq, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_join_metrics!(metrics, 0); + + // With batch coalescing, empty results may not emit any batches + // Check that either we have no batches, or an empty batch with proper schema + if batches_null_neq.is_empty() { + // This is fine - no output rows + } else { + let expected_null_neq = + ["+----+----+", "| n1 | n2 |", "+----+----+", "+----+----+"]; + assert_batches_eq!(expected_null_neq, &batches_null_neq); + } + + Ok(()) + } + + /// Returns the column names on the schema + fn columns(schema: &Schema) -> Vec<String> { + schema.fields().iter().map(|f| f.name().clone()).collect() + } + + /// This test verifies that the dynamic filter is marked as complete after HashJoinExec finishes building the hash table. + #[tokio::test] + async fn test_hash_join_marks_filter_complete() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 6]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (join, dynamic_filter) = + hash_join_with_dynamic_filter(left, right, on, JoinType::Inner)?; + + // Execute the join + let stream = join.execute(0, task_ctx)?; + let _batches = common::collect(stream).await?; + + // After the join completes, the dynamic filter should be marked as complete + // wait_complete() should return immediately + dynamic_filter.wait_complete().await; + + Ok(()) + } + + /// This test verifies that the dynamic filter is marked as complete even when the build side is empty. + #[tokio::test] + async fn test_hash_join_marks_filter_complete_empty_build_side() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + // Empty left side (build side) + let left = build_table(("a1", &vec![]), ("b1", &vec![]), ("c1", &vec![])); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (join, dynamic_filter) = + hash_join_with_dynamic_filter(left, right, on, JoinType::Inner)?; + + // Execute the join + let stream = join.execute(0, task_ctx)?; + let _batches = common::collect(stream).await?; + + // Even with empty build side, the dynamic filter should be marked as complete + // wait_complete() should return immediately + dynamic_filter.wait_complete().await; + + Ok(()) + } + + #[tokio::test] + async fn test_partitioned_dynamic_filter_reports_empty_canceled_partitions() + -> Result<()> { + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_dynamic_filter_pushdown = true; + let task_ctx = + Arc::new(TaskContext::default().with_session_config(session_config)); + + let child_left_schema = Arc::new(Schema::new(vec![ + Field::new("child_left_payload", DataType::Int32, false), + Field::new("child_key", DataType::Int32, false), + Field::new("child_left_extra", DataType::Int32, false), + ])); + let child_right_schema = Arc::new(Schema::new(vec![ + Field::new("child_right_payload", DataType::Int32, false), + Field::new("child_right_key", DataType::Int32, false), + Field::new("child_right_extra", DataType::Int32, false), + ])); + let parent_left_schema = Arc::new(Schema::new(vec![ + Field::new("parent_payload", DataType::Int32, false), + Field::new("parent_key", DataType::Int32, false), + Field::new("parent_extra", DataType::Int32, false), + ])); + + let child_left: Arc<dyn ExecutionPlan> = TestMemoryExec::try_new_exec( + &[ + vec![build_table_i32( + ("child_left_payload", &vec![10]), + ("child_key", &vec![0]), + ("child_left_extra", &vec![100]), + )], + vec![build_table_i32( + ("child_left_payload", &vec![11]), + ("child_key", &vec![1]), + ("child_left_extra", &vec![101]), + )], + vec![build_table_i32( + ("child_left_payload", &vec![12]), + ("child_key", &vec![2]), + ("child_left_extra", &vec![102]), + )], + vec![build_table_i32( + ("child_left_payload", &vec![13]), + ("child_key", &vec![3]), + ("child_left_extra", &vec![103]), + )], + ], + Arc::clone(&child_left_schema), + None, + )?; + let child_right: Arc<dyn ExecutionPlan> = TestMemoryExec::try_new_exec( + &[ + vec![build_table_i32( + ("child_right_payload", &vec![20]), + ("child_right_key", &vec![0]), + ("child_right_extra", &vec![200]), + )], + vec![build_table_i32( + ("child_right_payload", &vec![21]), + ("child_right_key", &vec![1]), + ("child_right_extra", &vec![201]), + )], + vec![build_table_i32( + ("child_right_payload", &vec![22]), + ("child_right_key", &vec![2]), + ("child_right_extra", &vec![202]), + )], + vec![build_table_i32( + ("child_right_payload", &vec![23]), + ("child_right_key", &vec![3]), + ("child_right_extra", &vec![203]), + )], + ], + Arc::clone(&child_right_schema), + None, + )?; + let parent_left: Arc<dyn ExecutionPlan> = TestMemoryExec::try_new_exec( + &[ + vec![build_table_i32( + ("parent_payload", &vec![30]), + ("parent_key", &vec![0]), + ("parent_extra", &vec![300]), + )], + vec![RecordBatch::new_empty(Arc::clone(&parent_left_schema))], + vec![build_table_i32( + ("parent_payload", &vec![32]), + ("parent_key", &vec![2]), + ("parent_extra", &vec![302]), + )], + vec![RecordBatch::new_empty(Arc::clone(&parent_left_schema))], + ], + Arc::clone(&parent_left_schema), + None, + )?; + + let child_on = vec![( + Arc::new(Column::new_with_schema("child_key", &child_left_schema)?) as _, + Arc::new(Column::new_with_schema( + "child_right_key", + &child_right_schema, + )?) as _, + )]; + let (child_join, _child_dynamic_filter) = hash_join_with_dynamic_filter_and_mode( + child_left, + child_right, + child_on, + JoinType::Inner, + PartitionMode::Partitioned, + )?; + let child_join: Arc<dyn ExecutionPlan> = Arc::new(child_join); + + let parent_on = vec![( + Arc::new(Column::new_with_schema("parent_key", &parent_left_schema)?) as _, + Arc::new(Column::new_with_schema("child_key", &child_join.schema())?) as _, + )]; + let parent_join = HashJoinExec::try_new( + parent_left, + child_join, + parent_on, + None, + &JoinType::RightSemi, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + + let batches = tokio::time::timeout( + std::time::Duration::from_secs(5), + crate::execution_plan::collect(Arc::new(parent_join), task_ctx), + ) + .await + .expect("partitioned right-semi join should not hang")?; + + assert_batches_sorted_eq!( + [ + "+--------------------+-----------+------------------+---------------------+-----------------+-------------------+", + "| child_left_payload | child_key | child_left_extra | child_right_payload | child_right_key | child_right_extra |", + "+--------------------+-----------+------------------+---------------------+-----------------+-------------------+", + "| 10 | 0 | 100 | 20 | 0 | 200 |", + "| 12 | 2 | 102 | 22 | 2 | 202 |", + "+--------------------+-----------+------------------+---------------------+-----------------+-------------------+", + ], + &batches + ); + + Ok(()) + } + + #[tokio::test] + async fn test_hash_join_skips_probe_on_empty_build_after_partition_bounds_report() + -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let (left, right, on) = empty_build_with_probe_error_inputs(); + + // Keep an extra consumer reference so execute() enables dynamic filter pushdown + // and enters the WaitPartitionBoundsReport path before deciding whether to poll + // the probe side. + let (join, dynamic_filter) = + hash_join_with_dynamic_filter(left, right, on, JoinType::Inner)?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + assert!(batches.is_empty()); + + dynamic_filter.wait_complete().await; + + Ok(()) + } + + #[tokio::test] + async fn test_perfect_hash_join_with_negative_numbers() -> Result<()> { + let task_ctx = prepare_task_ctx(8192, true); + let (left_schema, right_schema, on) = build_schema_and_on()?; + + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef, + Arc::new(Int32Array::from(vec![-1, 0, 1])) as ArrayRef, + ], + )?; + let left = TestMemoryExec::try_new_exec(&[vec![left_batch]], left_schema, None)?; + + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30, 40])) as ArrayRef, + Arc::new(Int32Array::from(vec![1, -1, 0, 2])) as ArrayRef, + ], + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None)?; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "a2", "b1"]); + + assert_batches_sorted_eq!( + [ + "+----+----+----+----+", + "| a1 | b1 | a2 | b1 |", + "+----+----+----+----+", + "| 1 | -1 | 20 | -1 |", + "| 2 | 0 | 30 | 0 |", + "| 3 | 1 | 10 | 1 |", + "+----+----+----+----+", + ], + &batches + ); + + assert_phj_used(&metrics, true); + + Ok(()) + } + + #[tokio::test] + async fn test_perfect_hash_join_overflow_full_int64_range() -> Result<()> { + let task_ctx = prepare_task_ctx(8192, true); + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![i64::MIN, i64::MAX]))], + )?; + let left = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::clone(&schema), + None, + )?; + let right = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)?; + let on: JoinOn = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let (_columns, batches, _metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 2); + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_phj_null_equals_null_build_no_nulls_probe_has_nulls( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let (left_schema, right_schema, on) = build_schema_and_on()?; + + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + Arc::new(Int32Array::from(vec![10, 20])) as ArrayRef, + ], + )?; + let left = TestMemoryExec::try_new_exec(&[vec![left_batch]], left_schema, None)?; + + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![3, 4])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(10), None])) as ArrayRef, + ], + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None)?; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNull, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "a2", "b1"]); + assert_batches_sorted_eq!( + [ + "+----+----+----+----+", + "| a1 | b1 | a2 | b1 |", + "+----+----+----+----+", + "| 1 | 10 | 3 | 10 |", + "+----+----+----+----+", + ], + &batches + ); + + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_phj_null_equals_nothing_build_probe_all_have_nulls( + batch_size: usize, + use_perfect_hash_join_as_possible: bool, + ) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, use_perfect_hash_join_as_possible); + let (left_schema, right_schema, on) = build_schema_and_on()?; + + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![Some(1), Some(2)])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(10), None])) as ArrayRef, + ], + )?; + let left = TestMemoryExec::try_new_exec(&[vec![left_batch]], left_schema, None)?; + + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![Some(3), Some(4)])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(10), None])) as ArrayRef, + ], + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None)?; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "a2", "b1"]); + assert_batches_sorted_eq!( + [ + "+----+----+----+----+", + "| a1 | b1 | a2 | b1 |", + "+----+----+----+----+", + "| 1 | 10 | 3 | 10 |", + "+----+----+----+----+", + ], + &batches + ); + + assert_phj_used(&metrics, use_perfect_hash_join_as_possible); + + Ok(()) + } + + #[tokio::test] + async fn test_phj_null_equals_null_build_have_nulls() -> Result<()> { + let task_ctx = prepare_task_ctx(8192, true); + let (left_schema, right_schema, on) = build_schema_and_on()?; + + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![Some(1), Some(2), Some(3)])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(10), Some(20), None])) as ArrayRef, + ], + )?; + let left = TestMemoryExec::try_new_exec(&[vec![left_batch]], left_schema, None)?; + + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![Some(3), Some(4)])) as ArrayRef, + Arc::new(Int32Array::from(vec![Some(10), Some(30)])) as ArrayRef, + ], + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None)?; + + let (columns, batches, metrics) = join_collect( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNull, + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "a2", "b1"]); + assert_batches_sorted_eq!( + [ + "+----+----+----+----+", + "| a1 | b1 | a2 | b1 |", + "+----+----+----+----+", + "| 1 | 10 | 3 | 10 |", + "+----+----+----+----+", + ], + &batches + ); + + assert_phj_used(&metrics, false); + + Ok(()) + } + + /// Test null-aware anti join when probe side (right) contains NULL + /// Expected: no rows should be output (NULL in subquery means all results are unknown) + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_anti_join_probe_null(batch_size: usize) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + + // Build left table (rows to potentially output) + let left = build_table_two_cols( + ("c1", &vec![Some(1), Some(2), Some(3), Some(4)]), + ("dummy", &vec![Some(10), Some(20), Some(30), Some(40)]), + ); + + // Build right table (subquery with NULL) + let right = build_table_two_cols( + ("c2", &vec![Some(1), Some(2), Some(3), None]), + ("dummy", &vec![Some(100), Some(200), Some(300), Some(400)]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("c2", &right.schema())?) as _, + )]; + + // Create null-aware anti join + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, // null_aware = true + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // Expected: empty result (probe side has NULL, so no rows should be output) + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + ++ + ++ + "); + } + Ok(()) + } + + /// Test null-aware anti join when build side (left) contains NULL keys + /// Expected: rows with NULL keys should not be output + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_anti_join_build_null(batch_size: usize) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + + // Build left table with NULL key (this row should not be output) + let left = build_table_two_cols( + ("c1", &vec![Some(1), Some(4), None]), + ("dummy", &vec![Some(10), Some(40), Some(0)]), + ); + + // Build right table (no NULL, so probe-side check passes) + let right = build_table_two_cols( + ("c2", &vec![Some(1), Some(2), Some(3)]), + ("dummy", &vec![Some(100), Some(200), Some(300)]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("c2", &right.schema())?) as _, + )]; + + // Create null-aware anti join + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, // null_aware = true + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // Expected: only c1=4 (not c1=1 which matches, not c1=NULL) + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+ + | c1 | dummy | + +----+-------+ + | 4 | 40 | + +----+-------+ + "); + } + Ok(()) + } + + /// Test null-aware anti join with no NULLs (should work like regular anti join) + #[apply(hash_join_exec_configs)] + #[tokio::test] + async fn test_null_aware_anti_join_no_nulls(batch_size: usize) -> Result<()> { + let task_ctx = prepare_task_ctx(batch_size, false); + + // Build left table (no NULLs) + let left = build_table_two_cols( + ("c1", &vec![Some(1), Some(2), Some(4), Some(5)]), + ("dummy", &vec![Some(10), Some(20), Some(40), Some(50)]), + ); + + // Build right table (no NULLs) + let right = build_table_two_cols( + ("c2", &vec![Some(1), Some(2), Some(3)]), + ("dummy", &vec![Some(100), Some(200), Some(300)]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("c2", &right.schema())?) as _, + )]; + + // Create null-aware anti join + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, // null_aware = true + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + // Expected: c1=4 and c1=5 (they don't match anything in right) + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+-------+ + | c1 | dummy | + +----+-------+ + | 4 | 40 | + | 5 | 50 | + +----+-------+ + "); + } + Ok(()) + } + + /// Test that null_aware validation rejects non-LeftAnti join types + #[tokio::test] + async fn test_null_aware_validation_wrong_join_type() { + let left = + build_table_two_cols(("c1", &vec![Some(1)]), ("dummy", &vec![Some(10)])); + let right = + build_table_two_cols(("c2", &vec![Some(1)]), ("dummy", &vec![Some(100)])); + + let on = vec![( + Arc::new(Column::new_with_schema("c1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("c2", &right.schema()).unwrap()) as _, + )]; + + // Try to create null-aware Inner join (should fail) + let result = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, // null_aware = true (invalid for Inner join) + ); + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("null_aware can only be true for LeftAnti joins") + ); + } + + /// Test that null_aware validation rejects multi-column joins + #[tokio::test] + async fn test_null_aware_validation_multi_column() { + let left = build_table(("a", &vec![1]), ("b", &vec![2]), ("c", &vec![3])); + let right = build_table(("x", &vec![1]), ("y", &vec![2]), ("z", &vec![3])); + + // Try multi-column join + let on = vec![ + ( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("x", &right.schema()).unwrap()) as _, + ), + ( + Arc::new(Column::new_with_schema("b", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("y", &right.schema()).unwrap()) as _, + ), + ]; + + // Try to create null-aware anti join with 2 columns (should fail) + let result = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, // null_aware = true (invalid for multi-column) + ); + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("null_aware anti join only supports single column join key") + ); + } + + #[test] + fn test_lr_is_preserved() { + assert_eq!(lr_is_preserved(JoinType::Inner), (true, true)); + assert_eq!(lr_is_preserved(JoinType::Left), (true, false)); + assert_eq!(lr_is_preserved(JoinType::Right), (false, true)); + assert_eq!(lr_is_preserved(JoinType::Full), (false, false)); + assert_eq!(lr_is_preserved(JoinType::LeftSemi), (true, true)); + assert_eq!(lr_is_preserved(JoinType::LeftAnti), (true, false)); + assert_eq!(lr_is_preserved(JoinType::LeftMark), (true, false)); + assert_eq!(lr_is_preserved(JoinType::RightSemi), (true, true)); + assert_eq!(lr_is_preserved(JoinType::RightAnti), (false, true)); + assert_eq!(lr_is_preserved(JoinType::RightMark), (false, true)); + } + + #[test] + fn test_with_dynamic_filter() -> Result<()> { + let (_, _, on) = build_schema_and_on()?; + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![1]), ("b1", &vec![1]), ("c2", &vec![1])); + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?; + assert!(join.dynamic_expressions_produced().is_empty()); + + let df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("b1", 1)) as _], + lit(true), + )); + let join = join.with_dynamic_filter_expr(Arc::clone(&df))?; + + let produced = join.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + assert_eq!( + produced[0] + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"), + df.expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"), + ); + Ok(()) + } + + #[test] + fn test_swap_inputs_rejects_dynamic_filter() -> Result<()> { + let left = build_table( + ("l_key", &vec![1]), + ("l_payload", &vec![10]), + ("l_other", &vec![100]), + ); + let right = build_table( + ("r_payload", &vec![20]), + ("r_key", &vec![1]), + ("r_other", &vec![200]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("l_key", &left.schema())?) as _, + Arc::new(Column::new_with_schema("r_key", &right.schema())?) as _, + )]; + + let dynamic_filter = HashJoinExec::create_dynamic_filter(&on); + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftSemi, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )? + .with_dynamic_filter_expr(dynamic_filter)?; + + let err = join.swap_inputs(PartitionMode::CollectLeft).unwrap_err(); + assert_contains!( + err.to_string(), + "Cannot swap HashJoinExec inputs after dynamic filters have been constructed" + ); + Ok(()) + } + + #[test] + fn test_dynamic_filter_pushdown_allowed_for_null_equal_join() -> Result<()> { + let (_, _, on) = build_schema_and_on()?; + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![1]), ("b1", &vec![1]), ("c2", &vec![1])); + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::RightSemi, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?; + + // Null-equal joins keep dynamic filter pushdown: the pushed predicate carries an + // `IS NULL` disjunct so a probe-side NULL still reaches the join. + assert!(join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_dynamic_filter_pushdown_rejects_null_aware_nullable_build_key() -> Result<()> + { + let left = build_table_two_cols( + ("a1", &vec![Some(1), None]), + ("b1", &vec![Some(1), Some(2)]), + ); + let right = build_table_two_cols( + ("a2", &vec![Some(2), Some(3)]), + ("b2", &vec![Some(1), Some(2)]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_dynamic_filter_pushdown_allows_null_aware_non_null_build_key() -> Result<()> { + // A NOT NULL build key cannot surface a build-side NULL, so the + // pushdown must stay enabled. + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![2]), ("b2", &vec![2]), ("c2", &vec![2])); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + assert!(join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + fn range_partitioned_dynamic_filter_test_join( + left_split: i32, + right_split: i32, + ) -> Result<(HashJoinExec, JoinOn)> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].0), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(left_split))])], + )?); + let right_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].1), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(right_split))])], + )?); + let left = Arc::new(PartitionedTestExec::try_new( + left_schema, + left_partitioning, + )?); + let right = Arc::new(PartitionedTestExec::try_new( + right_schema, + right_partitioning, + )?); + + let join = HashJoinExec::try_new( + left, + right, + on.clone(), + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + Ok((join, on)) + } + + fn with_hash_partitioned_children( + join: &HashJoinExec, + on: &JoinOn, + ) -> Result<HashJoinExec> { + join.builder() + .with_new_children(vec![ + Arc::new(PartitionedTestExec::try_new( + join.left().schema(), + Partitioning::Hash(vec![Arc::clone(&on[0].0)], 2), + )?), + Arc::new(PartitionedTestExec::try_new( + join.right().schema(), + Partitioning::Hash(vec![Arc::clone(&on[0].1)], 2), + )?), + ])? + .build() + } + + #[test] + fn test_partitioned_dynamic_filter_pushdown_allows_supported_partitioning() + -> Result<()> { + let (range_join, on) = range_partitioned_dynamic_filter_test_join(10, 10)?; + let hash_join = with_hash_partitioned_children(&range_join, &on)?; + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + assert!(range_join.allow_join_dynamic_filter_pushdown(session_config.options())); + assert!(hash_join.allow_join_dynamic_filter_pushdown(session_config.options())); + + session_config + .options_mut() + .optimizer + .preserve_file_partitions = 1; + assert!(range_join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_partitioned_dynamic_filter_pushdown_rejects_unsupported_partitioning() + -> Result<()> { + let (range_join, on) = range_partitioned_dynamic_filter_test_join(10, 10)?; + let hash_join = with_hash_partitioned_children(&range_join, &on)?; + let (mismatched_range_join, _) = + range_partitioned_dynamic_filter_test_join(10, 11)?; + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + assert!( + !mismatched_range_join + .allow_join_dynamic_filter_pushdown(session_config.options()) + ); + + session_config + .options_mut() + .optimizer + .preserve_file_partitions = 1; + assert!(!hash_join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { + let (_, _, on) = build_schema_and_on()?; + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![1]), ("b1", &vec![1]), ("c2", &vec![1])); + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?; + + // Column index 99 is out of bounds for the right (probe) side schema. + let df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("bad", 99)) as _], + lit(true), + )); + assert!(join.with_dynamic_filter_expr(df).is_err()); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/inlist_builder.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/inlist_builder.rs new file mode 100644 index 0000000..2fc3201 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/inlist_builder.rs
@@ -0,0 +1,158 @@ +// 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. + +//! Utilities for building InList expressions from hash join build side data + +use std::sync::Arc; + +use arrow::array::{ArrayRef, StructArray}; +use arrow::datatypes::{Field, FieldRef, Fields}; +use arrow_schema::DataType; +use datafusion_common::Result; + +pub(super) fn build_struct_fields(data_types: &[DataType]) -> Result<Fields> { + data_types + .iter() + .enumerate() + .map(|(i, dt)| Ok(Field::new(format!("c{i}"), dt.clone(), true))) + .collect() +} + +/// Builds InList values from join key column arrays. +/// +/// If `join_key_arrays` is: +/// 1. A single array, let's say Int32, this will produce a flat +/// InList expression where the lookup is expected to be scalar Int32 values, +/// that is: this will produce `IN LIST (1, 2, 3)` expected to be used as `2 IN LIST (1, 2, 3)`. +/// 2. An Int32 array and a Utf8 array, this will produce a Struct InList expression +/// where the lookup is expected to be Struct values with two fields (Int32, Utf8), +/// that is: this will produce `IN LIST ((1, "a"), (2, "b"))` expected to be used as `(2, "b") IN LIST ((1, "a"), (2, "b"))`. +/// The field names of the struct are auto-generated as "c0", "c1", ... and should match the struct expression used in the join keys. +/// +/// Note that this function does not deduplicate values - deduplication will happen later +/// when building an InList expression from this array via `InListExpr::try_new_from_array`. +/// +/// Returns `None` if the estimated size exceeds `max_size_bytes` or if the number of rows +/// exceeds `max_distinct_values`. +pub(super) fn build_struct_inlist_values( + join_key_arrays: &[ArrayRef], +) -> Result<Option<ArrayRef>> { + // Build the source array/struct + let source_array: ArrayRef = if join_key_arrays.len() == 1 { + // Single column: use directly + Arc::clone(&join_key_arrays[0]) + } else { + // Multi-column: build StructArray once from all columns + let fields = build_struct_fields( + &join_key_arrays + .iter() + .map(|arr| arr.data_type().clone()) + .collect::<Vec<_>>(), + )?; + + // Build field references with proper Arc wrapping + let arrays_with_fields: Vec<(FieldRef, ArrayRef)> = fields + .iter() + .cloned() + .zip(join_key_arrays.iter().cloned()) + .collect(); + + Arc::new(StructArray::from(arrays_with_fields)) + }; + + Ok(Some(source_array)) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + DictionaryArray, Int8Array, Int32Array, StringArray, StringDictionaryBuilder, + }; + + #[test] + fn test_build_single_column_inlist_array() { + let array = Arc::new(Int32Array::from(vec![1, 2, 3, 2, 1])) as ArrayRef; + let result = build_struct_inlist_values(std::slice::from_ref(&array)) + .unwrap() + .unwrap(); + + assert!(array.eq(&result)); + } + + #[test] + fn test_build_multi_column_inlist() { + let array1 = Arc::new(Int32Array::from(vec![1, 2, 3, 2, 1])) as ArrayRef; + let array2 = + Arc::new(StringArray::from(vec!["a", "b", "c", "b", "a"])) as ArrayRef; + + let result = build_struct_inlist_values(&[array1, array2]) + .unwrap() + .unwrap(); + + assert_eq!( + *result.data_type(), + DataType::Struct( + build_struct_fields(&[DataType::Int32, DataType::Utf8]).unwrap() + ) + ); + } + + #[test] + fn test_build_multi_column_inlist_with_dictionary() { + let mut builder = StringDictionaryBuilder::<arrow::datatypes::Int8Type>::new(); + builder.append_value("foo"); + builder.append_value("foo"); + builder.append_value("foo"); + let dict_array = Arc::new(builder.finish()) as ArrayRef; + + let int_array = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + + let result = build_struct_inlist_values(&[dict_array, int_array]) + .unwrap() + .unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!( + *result.data_type(), + DataType::Struct( + build_struct_fields(&[ + DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Utf8) + ), + DataType::Int32 + ]) + .unwrap() + ) + ); + } + + #[test] + fn test_build_single_column_dictionary_inlist() { + let keys = Int8Array::from(vec![0i8, 0, 0]); + let values = Arc::new(StringArray::from(vec!["foo"])); + let dict_array = Arc::new(DictionaryArray::new(keys, values)) as ArrayRef; + + let result = build_struct_inlist_values(std::slice::from_ref(&dict_array)) + .unwrap() + .unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!(result.data_type(), dict_array.data_type()); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs new file mode 100644 index 0000000..b915802 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/mod.rs
@@ -0,0 +1,27 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`HashJoinExec`] Partitioned Hash Join Operator + +pub use exec::{HashJoinExec, HashJoinExecBuilder}; +pub use partitioned_hash_eval::{HashExpr, HashTableLookupExpr, SeededRandomState}; + +mod exec; +mod inlist_builder; +mod partitioned_hash_eval; +mod shared_bounds; +mod stream;
diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/partitioned_hash_eval.rs new file mode 100644 index 0000000..08f23d9 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/partitioned_hash_eval.rs
@@ -0,0 +1,889 @@ +// 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. + +//! Hash computation and hash table lookup expressions for dynamic filtering + +use std::{fmt::Display, hash::Hash, sync::Arc}; + +use arrow::{ + array::{ArrayRef, UInt64Array}, + datatypes::{DataType, Schema}, + record_batch::RecordBatch, +}; +use datafusion_common::Result; +use datafusion_common::hash_utils::RandomState; +use datafusion_common::hash_utils::{create_hashes, with_hashes}; +#[cfg(feature = "proto")] +use datafusion_common::internal_err; +use datafusion_expr::ColumnarValue; +use datafusion_physical_expr_common::physical_expr::{ + DynHash, PhysicalExpr, PhysicalExprRef, +}; + +use crate::joins::Map; +use crate::joins::key_range_bitmap::KeyRangeBitmap; + +/// RandomState wrapper that preserves the seed used to create it. +/// +/// This is needed because `RandomState` doesn't expose its seed after creation, +/// but we need them for serialization (e.g., protobuf serde). +#[derive(Clone, Debug)] +pub struct SeededRandomState { + random_state: RandomState, + seed: u64, +} + +impl SeededRandomState { + /// Create a new SeededRandomState with the given seed. + pub const fn with_seed(k: u64) -> Self { + Self { + random_state: RandomState::with_seed(k), + seed: k, + } + } + + /// Get the inner RandomState. + pub fn random_state(&self) -> &RandomState { + &self.random_state + } + + /// Get the seed used to create this RandomState. + pub fn seed(&self) -> u64 { + self.seed + } +} + +/// Physical expression that computes hash values for a set of columns +/// +/// This expression computes the hash of join key columns using a specific RandomState. +/// It returns a UInt64Array containing the hash values. +/// +/// This is used for: +/// - Computing routing hashes (with RepartitionExec's 0,0,0,0 seeds) +/// - Computing lookup hashes (with HashJoin's 'J','O','I','N' seeds) +pub struct HashExpr { + /// Columns to hash + on_columns: Vec<PhysicalExprRef>, + /// Random state for hashing (with seeds preserved for serialization) + random_state: SeededRandomState, + /// Description for display + description: String, +} + +impl HashExpr { + /// Create a new HashExpr + /// + /// # Arguments + /// * `on_columns` - Columns to hash + /// * `random_state` - SeededRandomState for hashing + /// * `description` - Description for debugging (e.g., "hash_repartition", "hash_join") + pub fn new( + on_columns: Vec<PhysicalExprRef>, + random_state: SeededRandomState, + description: String, + ) -> Self { + Self { + on_columns, + random_state, + description, + } + } + + /// Get the columns being hashed. + pub fn on_columns(&self) -> &[PhysicalExprRef] { + &self.on_columns + } + + /// Get the seed used for hashing. + pub fn seed(&self) -> u64 { + self.random_state.seed() + } + + /// Get the description. + pub fn description(&self) -> &str { + &self.description + } +} + +impl std::fmt::Debug for HashExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let cols = self + .on_columns + .iter() + .map(|e| e.to_string()) + .collect::<Vec<_>>() + .join(", "); + let seed = self.seed(); + write!(f, "{}({cols}, [{seed}])", self.description) + } +} + +impl Hash for HashExpr { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.on_columns.dyn_hash(state); + self.description.hash(state); + self.seed().hash(state); + } +} + +impl PartialEq for HashExpr { + fn eq(&self, other: &Self) -> bool { + self.on_columns == other.on_columns + && self.description == other.description + && self.seed() == other.seed() + } +} + +impl Eq for HashExpr {} + +impl Display for HashExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.description) + } +} + +impl PhysicalExpr for HashExpr { + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + self.on_columns.iter().collect() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn PhysicalExpr>>, + ) -> Result<Arc<dyn PhysicalExpr>> { + Ok(Arc::new(HashExpr::new( + children, + self.random_state.clone(), + self.description.clone(), + ))) + } + + fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { + Ok(DataType::UInt64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result<bool> { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + let num_rows = batch.num_rows(); + + // Evaluate columns + let keys_values = evaluate_columns(&self.on_columns, batch)?; + + // Compute hashes + let mut hashes_buffer = vec![0; num_rows]; + create_hashes( + &keys_values, + self.random_state.random_state(), + &mut hashes_buffer, + )?; + + Ok(ColumnarValue::Array(Arc::new(UInt64Array::from( + hashes_buffer, + )))) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.description) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> { + use datafusion_proto_models::protobuf; + let on_columns = ctx.encode_children_expressions(&self.on_columns)?; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( + protobuf::PhysicalHashExprNode { + on_columns, + seed0: self.seed(), + description: self.description.clone(), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl HashExpr { + /// Reconstruct a [`HashExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`], the exact inverse of what + /// [`PhysicalExpr::try_to_proto`] produces, so every expression's + /// `try_from_proto` shares one signature. Child sub-expressions are + /// decoded recursively via [`PhysicalExprDecodeCtx::decode`]. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto + /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result<Arc<dyn PhysicalExpr>> { + use datafusion_proto_models::protobuf; + let hash_expr = match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::HashExpr(h)) => h, + _ => return internal_err!("PhysicalExprNode is not a HashExpr"), + }; + let on_columns = ctx.decode_children_expressions(&hash_expr.on_columns)?; + Ok(Arc::new(HashExpr::new( + on_columns, + SeededRandomState::with_seed(hash_expr.seed0), + hash_expr.description.clone(), + ))) + } +} + +/// Physical expression that checks join keys in a [`Map`] (hash table or array map). +/// +/// Returns a [`BooleanArray`](arrow::array::BooleanArray) indicating if join keys (from `on_columns`) exist in the map. +// TODO: rename to MapLookupExpr +pub struct HashTableLookupExpr { + /// Columns in the ON clause used to compute the join key for lookups + on_columns: Vec<PhysicalExprRef>, + /// Random state for hashing (with seeds preserved for serialization) + random_state: SeededRandomState, + /// Map to check against (hash table or array map) + map: Arc<Map>, + /// Description for display + description: String, + /// Shared pruning-only view of the build side's key range + pruning_bitmap: Option<Arc<KeyRangeBitmap>>, +} +impl HashTableLookupExpr { + /// Create a new HashTableLookupExpr + /// + /// # Arguments + /// * `on_columns` - Columns in the ON clause used to compute the join key + /// * `random_state` - SeededRandomState for hashing + /// * `map` - Map to check membership (hash table or array map) + /// * `description` - Description for debugging + /// # Note + /// This is public for internal testing purposes only and is not + /// guaranteed to be stable across versions. + pub fn new( + on_columns: Vec<PhysicalExprRef>, + random_state: SeededRandomState, + map: Arc<Map>, + description: String, + ) -> Self { + Self { + on_columns, + random_state, + map, + description, + pruning_bitmap: None, + } + } + + /// Attach a key-range summary for container pruning. + pub fn with_pruning_bitmap(mut self, bitmap: Option<Arc<KeyRangeBitmap>>) -> Self { + self.pruning_bitmap = bitmap; + self + } + + /// Which parts of this build side's key range hold keys, for container pruning. + pub fn pruning_bitmap(&self) -> Option<&Arc<KeyRangeBitmap>> { + self.pruning_bitmap.as_ref() + } +} +impl std::fmt::Debug for HashTableLookupExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let cols = self + .on_columns + .iter() + .map(|e| e.to_string()) + .collect::<Vec<_>>() + .join(", "); + let seed = self.random_state.seed(); + write!(f, "{}({cols}, [{seed}])", self.description) + } +} + +impl Hash for HashTableLookupExpr { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.on_columns.dyn_hash(state); + self.description.hash(state); + self.random_state.seed().hash(state); + // Note that we compare hash_map by pointer equality. + // Actually comparing the contents of the hash maps would be expensive. + // The way these hash maps are used in actuality is that HashJoinExec creates + // one per partition per query execution, thus it is never possible for two different + // hash maps to have the same content in practice. + // Theoretically this is a public API and users could create identical hash maps, + // but that seems unlikely and not worth paying the cost of deep comparison all the time. + Arc::as_ptr(&self.map).hash(state); + } +} + +impl PartialEq for HashTableLookupExpr { + fn eq(&self, other: &Self) -> bool { + // Note that we compare hash_map by pointer equality. + // Actually comparing the contents of the hash maps would be expensive. + // The way these hash maps are used in actuality is that HashJoinExec creates + // one per partition per query execution, thus it is never possible for two different + // hash maps to have the same content in practice. + // Theoretically this is a public API and users could create identical hash maps, + // but that seems unlikely and not worth paying the cost of deep comparison all the time. + self.on_columns == other.on_columns + && self.description == other.description + && self.random_state.seed() == other.random_state.seed() + && Arc::ptr_eq(&self.map, &other.map) + } +} + +impl Eq for HashTableLookupExpr {} + +impl Display for HashTableLookupExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.description) + } +} + +impl PhysicalExpr for HashTableLookupExpr { + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + self.on_columns.iter().collect() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn PhysicalExpr>>, + ) -> Result<Arc<dyn PhysicalExpr>> { + Ok(Arc::new(HashTableLookupExpr { + on_columns: children, + random_state: self.random_state.clone(), + map: Arc::clone(&self.map), + description: self.description.clone(), + pruning_bitmap: self.pruning_bitmap.clone(), + })) + } + + fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result<bool> { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + // Evaluate columns + let join_keys = evaluate_columns(&self.on_columns, batch)?; + + match self.map.as_ref() { + Map::HashMap(map) => { + with_hashes(&join_keys, self.random_state.random_state(), |hashes| { + let array = map.contain_hashes(hashes); + Ok(ColumnarValue::Array(Arc::new(array))) + }) + } + Map::ArrayMap(map) => { + let array = map.contain_keys(&join_keys)?; + Ok(ColumnarValue::Array(Arc::new(array))) + } + } + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> { + use datafusion_proto_models::protobuf; + use datafusion_proto_models::protobuf::physical_expr_node::ExprType; + + // HashTableLookupExpr holds a runtime Arc<Map> (the build-side hash + // table) that cannot be serialized, so it is replaced with lit(true). + // + // Dynamic filtering is a performance optimisation only — replacing the + // lookup with lit(true) preserves correctness by allowing all rows + // through. + // + // If a plan is serialized before execution, HashTableLookupExpr is not + // yet present in the dynamic filter expression. + // + // If a plan is serialized after execution, any runtime-created + // HashTableLookupExpr is replaced during serialization. Re-executing + // the plan requires reset_state(), after which HashJoinExec rebuilds + // fresh dynamic filters at runtime. + let value = datafusion_proto_common::ScalarValue { + value: Some(datafusion_proto_common::scalar_value::Value::BoolValue( + true, + )), + }; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(ExprType::Literal(value)), + })) + } + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.description) + } +} + +fn evaluate_columns( + columns: &[PhysicalExprRef], + batch: &RecordBatch, +) -> Result<Vec<ArrayRef>> { + let num_rows = batch.num_rows(); + columns + .iter() + .map(|c| c.evaluate(batch)?.into_array(num_rows)) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::joins::join_hash_map::JoinHashMapU32; + use datafusion_common::ScalarValue; + use datafusion_physical_expr::expressions::Column; + use std::collections::hash_map::DefaultHasher; + use std::hash::Hasher; + + fn compute_hash<T: Hash>(value: &T) -> u64 { + let mut hasher = DefaultHasher::new(); + value.hash(&mut hasher); + hasher.finish() + } + + /// The pruning bitmap must survive `with_new_children`, which a scan calls + /// once per file. + #[test] + fn test_pruning_bitmap_shared_with_derived_children() { + let keys: ArrayRef = Arc::new(arrow::array::Int64Array::from(vec![0, 1_000_000])); + let bitmap = Arc::new( + KeyRangeBitmap::try_new( + &keys, + &ScalarValue::Int64(Some(0)), + &ScalarValue::Int64(Some(1_000_000)), + 2, + ) + .expect("two keys a million apart leave prunable gaps"), + ); + let expr = Arc::new( + HashTableLookupExpr::new( + vec![Arc::new(Column::new("a", 0))], + SeededRandomState::with_seed(1), + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))), + "hash_lookup".to_string(), + ) + .with_pruning_bitmap(Some(Arc::clone(&bitmap))), + ); + + let derived = expr + .with_new_children(vec![Arc::new(Column::new("a", 7))]) + .unwrap(); + let derived = derived.downcast_ref::<HashTableLookupExpr>().unwrap(); + assert_eq!(derived.children()[0].to_string(), "a@7"); + assert!(Arc::ptr_eq(&bitmap, derived.pruning_bitmap().unwrap())); + } + + #[test] + fn test_hash_expr_eq_same() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let col_b: PhysicalExprRef = Arc::new(Column::new("b", 1)); + + let expr1 = HashExpr::new( + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + let expr2 = HashExpr::new( + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + assert_eq!(expr1, expr2); + } + + #[test] + fn test_hash_expr_eq_different_columns() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let col_b: PhysicalExprRef = Arc::new(Column::new("b", 1)); + let col_c: PhysicalExprRef = Arc::new(Column::new("c", 2)); + + let expr1 = HashExpr::new( + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + let expr2 = HashExpr::new( + vec![Arc::clone(&col_a), Arc::clone(&col_c)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + assert_ne!(expr1, expr2); + } + + #[test] + fn test_hash_expr_eq_different_description() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + + let expr1 = HashExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + "hash_one".to_string(), + ); + + let expr2 = HashExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + "hash_two".to_string(), + ); + + assert_ne!(expr1, expr2); + } + + #[test] + fn test_hash_expr_eq_different_seeds() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + + let expr1 = HashExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + let expr2 = HashExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(5), + "test_hash".to_string(), + ); + + assert_ne!(expr1, expr2); + } + + #[test] + fn test_hash_expr_hash_consistency() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let col_b: PhysicalExprRef = Arc::new(Column::new("b", 1)); + + let expr1 = HashExpr::new( + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + let expr2 = HashExpr::new( + vec![Arc::clone(&col_a), Arc::clone(&col_b)], + SeededRandomState::with_seed(1), + "test_hash".to_string(), + ); + + // Equal expressions should have equal hashes + assert_eq!(expr1, expr2); + assert_eq!(compute_hash(&expr1), compute_hash(&expr2)); + } + + #[cfg(feature = "proto")] + mod proto_tests { + use super::*; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, + }; + use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, + }; + use datafusion_proto_models::protobuf; + + struct TestEncoder; + + impl PhysicalExprEncode for TestEncoder { + fn encode( + &self, + expr: &Arc<dyn PhysicalExpr>, + ) -> Result<protobuf::PhysicalExprNode> { + let ctx = PhysicalExprEncodeCtx::new(self); + expr.try_to_proto(&ctx)?.ok_or_else(|| { + internal_datafusion_err!("test encoder cannot encode {expr:?}") + }) + } + } + + struct TestDecoder; + + impl PhysicalExprDecode for TestDecoder { + fn decode( + &self, + node: &protobuf::PhysicalExprNode, + schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>> { + let ctx = PhysicalExprDecodeCtx::new(schema, self); + match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::Column(_)) => { + Column::try_from_proto(node, &ctx) + } + _ => internal_err!("test decoder cannot decode {node:?}"), + } + } + } + + fn test_decode_ctx<'a>( + schema: &'a Schema, + decoder: &'a TestDecoder, + ) -> PhysicalExprDecodeCtx<'a> { + PhysicalExprDecodeCtx::new(schema, decoder) + } + + #[test] + fn hash_expr_try_to_proto() { + let expr = HashExpr::new( + vec![Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1))], + SeededRandomState::with_seed(42), + "hash_join".to_string(), + ); + let encoder = TestEncoder; + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let proto = expr.try_to_proto(&ctx).unwrap().unwrap(); + + assert_eq!(proto.expr_id, None); + let hash_expr = match proto.expr_type.unwrap() { + protobuf::physical_expr_node::ExprType::HashExpr(hash_expr) => hash_expr, + other => panic!("expected HashExpr, got {other:?}"), + }; + assert_eq!(hash_expr.seed0, 42); + assert_eq!(hash_expr.description, "hash_join"); + assert_eq!(hash_expr.on_columns.len(), 2); + assert!( + hash_expr + .on_columns + .iter() + .all(|expr| expr.expr_id.is_none()) + ); + } + + #[test] + fn hash_expr_try_from_proto() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ]); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( + protobuf::PhysicalHashExprNode { + on_columns: vec![ + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some( + protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: "a".to_string(), + index: 0, + }, + ), + ), + }, + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some( + protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: "b".to_string(), + index: 1, + }, + ), + ), + }, + ], + seed0: 42, + description: "hash_join".to_string(), + }, + )), + }; + + let expr = HashExpr::try_from_proto(&proto, &ctx).unwrap(); + let expr = expr.downcast_ref::<HashExpr>().unwrap(); + + assert_eq!(expr.seed(), 42); + assert_eq!(expr.description(), "hash_join"); + assert_eq!(expr.on_columns().len(), 2); + assert_eq!( + expr.on_columns()[0] + .downcast_ref::<Column>() + .map(|col| (col.name(), col.index())), + Some(("a", 0)) + ); + assert_eq!( + expr.on_columns()[1] + .downcast_ref::<Column>() + .map(|col| (col.name(), col.index())), + Some(("b", 1)) + ); + } + + #[test] + fn hash_expr_try_from_proto_rejects_wrong_node_type() { + let schema = Schema::empty(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: "a".to_string(), + index: 0, + }, + )), + }; + + let err = HashExpr::try_from_proto(&proto, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalExprNode is not a HashExpr"), + "{err}" + ); + } + } + + #[test] + fn test_hash_table_lookup_expr_eq_same() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let hash_map = + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))); + + let expr1 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup".to_string(), + ); + + let expr2 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup".to_string(), + ); + + assert_eq!(expr1, expr2); + } + + #[test] + fn test_hash_table_lookup_expr_eq_different_columns() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let col_b: PhysicalExprRef = Arc::new(Column::new("b", 1)); + + let hash_map = + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))); + + let expr1 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup".to_string(), + ); + + let expr2 = HashTableLookupExpr::new( + vec![Arc::clone(&col_b)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup".to_string(), + ); + + assert_ne!(expr1, expr2); + } + + #[test] + fn test_hash_table_lookup_expr_eq_different_description() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let hash_map = + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))); + + let expr1 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup_one".to_string(), + ); + + let expr2 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup_two".to_string(), + ); + + assert_ne!(expr1, expr2); + } + + #[test] + fn test_hash_table_lookup_expr_eq_different_hash_map() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + + // Two different Arc pointers (even with same content) should not be equal + let hash_map1 = + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))); + let hash_map2 = + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))); + let expr1 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + hash_map1, + "lookup".to_string(), + ); + + let expr2 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + hash_map2, + "lookup".to_string(), + ); + + // Different Arc pointers means not equal (uses Arc::ptr_eq) + assert_ne!(expr1, expr2); + } + + #[test] + fn test_hash_table_lookup_expr_hash_consistency() { + let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); + let hash_map = + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(10)))); + + let expr1 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup".to_string(), + ); + + let expr2 = HashTableLookupExpr::new( + vec![Arc::clone(&col_a)], + SeededRandomState::with_seed(1), + Arc::clone(&hash_map), + "lookup".to_string(), + ); + + // Equal expressions should have equal hashes + assert_eq!(expr1, expr2); + assert_eq!(compute_hash(&expr1), compute_hash(&expr2)); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/shared_bounds.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/shared_bounds.rs new file mode 100644 index 0000000..3ecebd2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/shared_bounds.rs
@@ -0,0 +1,1522 @@ +// 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. + +//! Utilities for shared build-side information. Used in dynamic filter pushdown in Hash Joins. +// TODO: include the link to the Dynamic Filter blog post. + +use std::fmt; +use std::sync::Arc; + +use crate::ExecutionPlan; +use crate::ExecutionPlanProperties; +use crate::Partitioning; +use crate::joins::Map; +use crate::joins::PartitionMode; +use crate::joins::hash_join::exec::HASH_JOIN_SEED; +use crate::joins::hash_join::inlist_builder::build_struct_fields; +use crate::joins::hash_join::partitioned_hash_eval::{ + HashExpr, HashTableLookupExpr, SeededRandomState, +}; +use crate::joins::key_range_bitmap::KeyRangeBitmap; +use crate::repartition::RangeExpr; +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{ + DataFusionError, NullEquality, Result, ScalarValue, SharedResult, + assert_or_internal_err, +}; +use datafusion_expr::Operator; +use datafusion_functions::core::r#struct as struct_func; +use datafusion_physical_expr::expressions::{ + BinaryExpr, CaseExpr, DynamicFilterPhysicalExpr, InListExpr, IsNullExpr, lit, +}; +use datafusion_physical_expr::{ + PhysicalExpr, PhysicalExprRef, RangePartitioning, ScalarFunctionExpr, +}; + +use parking_lot::Mutex; +use tokio::sync::Notify; + +/// Represents the minimum and maximum values for a specific column. +/// Used in dynamic filter pushdown to establish value boundaries. +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ColumnBounds { + /// The minimum value observed for this column + pub(crate) min: ScalarValue, + /// The maximum value observed for this column + pub(crate) max: ScalarValue, +} + +impl ColumnBounds { + pub(crate) fn new(min: ScalarValue, max: ScalarValue) -> Self { + Self { min, max } + } +} + +/// Represents the bounds for all join key columns from a single partition. +/// This contains the min/max values computed from one partition's build-side data. +#[derive(Debug, Clone)] +pub(crate) struct PartitionBounds { + /// Min/max bounds for each join key column in this partition. + /// Index corresponds to the join key expression index. + column_bounds: Vec<ColumnBounds>, +} + +impl PartitionBounds { + pub(crate) fn new(column_bounds: Vec<ColumnBounds>) -> Self { + Self { column_bounds } + } + + pub(crate) fn get_column_bounds(&self, index: usize) -> Option<&ColumnBounds> { + self.column_bounds.get(index) + } +} + +/// Creates a membership predicate for filter pushdown. +/// +/// If `inlist_values` is provided (for small build sides), creates an InList expression. +/// Otherwise, creates a HashTableLookup expression (for large build sides). +/// +/// Supports both single-column and multi-column joins using struct expressions. +fn create_membership_predicate( + on_right: &[PhysicalExprRef], + pushdown: PushdownStrategy, + random_state: &SeededRandomState, + schema: &Schema, +) -> Result<Option<Arc<dyn PhysicalExpr>>> { + match pushdown { + // Use InList expression for small build sides + PushdownStrategy::InList(in_list_array) => { + // Build the expression to compare against + let expr = if on_right.len() == 1 { + // Single column: col IN (val1, val2, ...) + Arc::clone(&on_right[0]) + } else { + let fields = build_struct_fields( + on_right + .iter() + .map(|r| r.data_type(schema)) + .collect::<Result<Vec<_>>>()? + .as_ref(), + )?; + + // The return field name and the function field name don't really matter here. + let return_field = + Arc::new(Field::new("struct", DataType::Struct(fields), true)); + + Arc::new(ScalarFunctionExpr::new( + "struct", + struct_func(), + on_right.to_vec(), + return_field, + Arc::new(ConfigOptions::default()), + )) as Arc<dyn PhysicalExpr> + }; + + // Use InListExpr::try_new_from_array() to build an InList with static_filter optimization (hash-based lookup) + Ok(Some(Arc::new(InListExpr::try_new_from_array( + expr, + in_list_array, + false, + schema, + )?))) + } + // Use hash table lookup for large build sides + PushdownStrategy::Map(hash_map, pruning_bitmap) => Ok(Some(Arc::new( + HashTableLookupExpr::new( + on_right.to_vec(), + random_state.clone(), + hash_map, + "hash_lookup".to_string(), + ) + .with_pruning_bitmap(pruning_bitmap), + ) + as Arc<dyn PhysicalExpr>)), + // Empty partition - should not create a filter for this + PushdownStrategy::Empty => Ok(None), + } +} + +/// Creates a bounds predicate from partition bounds. +/// +/// Returns `None` if no column bounds are available. +/// Returns a combined predicate (col >= min AND col <= max) for all columns with bounds. +fn create_bounds_predicate( + on_right: &[PhysicalExprRef], + bounds: &PartitionBounds, +) -> Option<Arc<dyn PhysicalExpr>> { + let mut column_predicates = Vec::new(); + + for (col_idx, right_expr) in on_right.iter().enumerate() { + if let Some(column_bounds) = bounds.get_column_bounds(col_idx) { + // Create predicate: col >= min AND col <= max + let min_expr = Arc::new(BinaryExpr::new( + Arc::clone(right_expr), + Operator::GtEq, + lit(column_bounds.min.clone()), + )) as Arc<dyn PhysicalExpr>; + let max_expr = Arc::new(BinaryExpr::new( + Arc::clone(right_expr), + Operator::LtEq, + lit(column_bounds.max.clone()), + )) as Arc<dyn PhysicalExpr>; + let range_expr = Arc::new(BinaryExpr::new(min_expr, Operator::And, max_expr)) + as Arc<dyn PhysicalExpr>; + column_predicates.push(range_expr); + } + } + + if column_predicates.is_empty() { + None + } else { + Some( + column_predicates + .into_iter() + .reduce(|acc, pred| { + Arc::new(BinaryExpr::new(acc, Operator::And, pred)) + as Arc<dyn PhysicalExpr> + }) + .unwrap(), + ) + } +} + +/// Combines a membership predicate and a bounds predicate with logical AND. +/// +/// Returns `None` when neither is available; callers decide the fallback (e.g. +/// skip updating the filter vs. emit a `lit(true)` branch inside a CASE). +fn combine_membership_and_bounds( + membership_expr: Option<Arc<dyn PhysicalExpr>>, + bounds_expr: Option<Arc<dyn PhysicalExpr>>, +) -> Option<Arc<dyn PhysicalExpr>> { + match (membership_expr, bounds_expr) { + (Some(membership), Some(bounds)) => { + Some(Arc::new(BinaryExpr::new(bounds, Operator::And, membership)) + as Arc<dyn PhysicalExpr>) + } + (Some(membership), None) => Some(membership), + (None, Some(bounds)) => Some(bounds), + (None, None) => None, + } +} + +/// Coordinates build-side information collection across multiple partitions +/// +/// This structure collects information from the build side (hash tables and/or bounds) and +/// ensures that dynamic filters are built with complete information from all relevant +/// partitions before being applied to probe-side scans. Incomplete filters would +/// incorrectly eliminate valid join results. +/// +/// ## Synchronization Strategy +/// +/// 1. Each partition computes information from its build-side data (hash maps and/or bounds) +/// 2. Information is stored in the shared state, which tracks how many partitions have reported +/// 3. When the last partition reports, one waiter is elected as the finalizer; it merges the +/// collected information, updates the dynamic filter exactly once, and publishes the +/// terminal result by transitioning [`CompletionState`] to `Ready` +/// 4. A [`tokio::sync::Notify`] wakes any other partitions parked in `wait_for_completion`, +/// which then observe the `Ready` state under the mutex and return immediately +/// +/// ## Hash Map vs Bounds +/// +/// - **Hash Maps (Partitioned mode)**: Collects Arc references to hash tables from each partition. +/// Creates a `PartitionedHashLookupPhysicalExpr` that routes rows to the correct partition's hash table. +/// - **Bounds (CollectLeft mode)**: Collects min/max bounds and creates range predicates. +/// +/// ## Partition Counting +/// +/// The `total_partitions` count represents how many times `collect_build_side` will be called: +/// - **CollectLeft**: Number of output partitions (each accesses shared build data) +/// - **Partitioned**: Number of input partitions (each builds independently) +/// +/// ## Thread Safety +/// +/// All fields use a single mutex to ensure correct coordination between concurrent +/// partition executions. +pub(crate) struct SharedBuildAccumulator { + /// Build-side data protected by a single mutex to avoid ordering concerns + inner: Mutex<AccumulatorState>, + /// Wakes every partition that is parked in [`Self::wait_for_completion`] + /// once [`AccumulatorState::completion`] transitions to + /// [`CompletionState::Ready`]. Notifications are fired once per + /// accumulator lifetime (the elected finalizer publishes the terminal + /// result, then broadcasts), so late subscribers simply re-check the + /// state under the mutex and return immediately. + completion_notify: Notify, + /// Dynamic filter for pushdown to probe side + dynamic_filter: Arc<DynamicFilterPhysicalExpr>, + /// Right side join expressions needed for creating filter expressions + on_right: Vec<PhysicalExprRef>, + /// Random state for partitioning (RepartitionExec's hash function with 0,0,0,0 seeds) + /// Used for PartitionedHashLookupPhysicalExpr + repartition_random_state: SeededRandomState, + /// Schema of the probe (right) side for evaluating filter expressions + probe_schema: Arc<Schema>, + /// Probe-side Range routing metadata for partitioned dynamic filters. + probe_range_partitioning: Option<RangePartitioning>, + /// Null equality of the join. Under `NullEqualsNull` a probe-side NULL can match a + /// build-side NULL, so the pushed filter must keep NULL rows here too. + null_equality: NullEquality, + /// Null-aware anti join (`NOT IN`). A probe-side NULL must reach the join so its + /// three-valued logic can collapse the result, so the pushed filter keeps NULL rows. + null_aware: bool, +} + +/// Strategy for filter pushdown (decided at collection time) +#[derive(Clone)] +pub(crate) enum PushdownStrategy { + /// Use InList for small build sides (< 128MB) + InList(ArrayRef), + /// Use map lookup for large build sides. If `Some` the second field + /// represents the build side's keys using a bucket bitmap. + Map(Arc<Map>, Option<Arc<KeyRangeBitmap>>), + /// There was no data in this partition, do not build a dynamic filter for it + Empty, +} + +/// Build-side data reported by a single partition +pub(crate) enum PartitionBuildData { + Partitioned { + partition_id: usize, + pushdown: PushdownStrategy, + bounds: PartitionBounds, + keys_have_null: bool, + }, + CollectLeft { + pushdown: PushdownStrategy, + bounds: PartitionBounds, + keys_have_null: bool, + }, +} + +/// Per-partition accumulated data (Partitioned mode) +#[derive(Clone)] +struct PartitionData { + bounds: PartitionBounds, + pushdown: PushdownStrategy, + /// Whether any build key of this partition is NULL. Decides whether the pushed + /// filter must keep probe-side NULL rows for a null-equal join to match them. + keys_have_null: bool, +} + +/// Build-side data organized by partition mode +enum AccumulatedBuildData { + Partitioned { + partitions: Vec<PartitionStatus>, + completed_partitions: usize, + }, + CollectLeft { + data: PartitionStatus, + reported_count: usize, + expected_reports: usize, + }, +} + +enum CompletionState { + Pending, + Finalizing, + Ready(SharedResult<()>), +} + +struct AccumulatorState { + data: AccumulatedBuildData, + completion: CompletionState, +} + +#[derive(Clone)] +enum PartitionStatus { + Pending, + Reported(PartitionData), + CanceledUnknown, +} + +#[derive(Clone)] +enum FinalizeInput { + Partitioned(Vec<PartitionStatus>), + CollectLeft(PartitionStatus), +} + +impl SharedBuildAccumulator { + /// Creates a new SharedBuildAccumulator configured for the given partition mode + /// + /// This method calculates how many times `collect_build_side` will be called based on the + /// partition mode's execution pattern. This count is critical for determining when we have + /// complete information from all partitions to build the dynamic filter. + /// + /// ## Partition Mode Execution Patterns + /// + /// - **CollectLeft**: Build side is collected ONCE from partition 0 and shared via `OnceFut` + /// across all output partitions. Each output partition calls `collect_build_side` to access the shared build data. + /// Although this results in multiple invocations, the `report_partition_bounds` function contains deduplication logic to handle them safely. + /// Expected calls = number of output partitions. + /// + /// + /// - **Partitioned**: Each partition independently builds its own hash table by calling + /// `collect_build_side` once. Expected calls = number of build partitions. + /// + /// - **Auto**: Placeholder mode resolved during optimization. Uses 1 as safe default since + /// the actual mode will be determined and a new accumulator created before execution. + /// + /// ## Why This Matters + /// + /// We cannot build a partial filter from some partitions - it would incorrectly eliminate + /// valid join results. We must wait until we have complete information from ALL + /// relevant partitions before updating the dynamic filter. + #[expect(clippy::too_many_arguments)] + pub(crate) fn new_from_partition_mode( + partition_mode: PartitionMode, + left_child: &dyn ExecutionPlan, + right_child: &dyn ExecutionPlan, + dynamic_filter: Arc<DynamicFilterPhysicalExpr>, + on_right: Vec<PhysicalExprRef>, + repartition_random_state: SeededRandomState, + null_equality: NullEquality, + null_aware: bool, + ) -> Self { + // Troubleshooting: If partition counts are incorrect, verify this logic matches + // the actual execution pattern in collect_build_side() + let expected_calls = match partition_mode { + // Each output partition accesses shared build data + PartitionMode::CollectLeft => { + right_child.output_partitioning().partition_count() + } + // Each partition builds its own data + PartitionMode::Partitioned => { + left_child.output_partitioning().partition_count() + } + // Default value, will be resolved during optimization (does not exist once `execute()` is called; will be replaced by one of the other two) + PartitionMode::Auto => unreachable!( + "PartitionMode::Auto should not be present at execution time. This is a bug in DataFusion, please report it!" + ), + }; + + let mode_data = match partition_mode { + PartitionMode::Partitioned => AccumulatedBuildData::Partitioned { + partitions: vec![ + PartitionStatus::Pending; + left_child.output_partitioning().partition_count() + ], + completed_partitions: 0, + }, + PartitionMode::CollectLeft => AccumulatedBuildData::CollectLeft { + data: PartitionStatus::Pending, + reported_count: 0, + expected_reports: expected_calls, + }, + PartitionMode::Auto => unreachable!( + "PartitionMode::Auto should not be present at execution time. This is a bug in DataFusion, please report it!" + ), + }; + + let probe_range_partitioning = + match (partition_mode, right_child.output_partitioning()) { + (PartitionMode::Partitioned, Partitioning::Range(range)) => { + Some(range.clone()) + } + _ => None, + }; + + Self { + inner: Mutex::new(AccumulatorState { + data: mode_data, + completion: CompletionState::Pending, + }), + completion_notify: Notify::new(), + dynamic_filter, + on_right, + repartition_random_state, + probe_schema: right_child.schema(), + probe_range_partitioning, + null_equality, + null_aware, + } + } + + /// Report build-side data from a partition + /// + /// This unified method handles both CollectLeft and Partitioned modes. When all partitions + /// have reported (barrier wait), the leader builds the appropriate filter expression: + /// - CollectLeft: Simple conjunction of bounds and membership check + /// - Partitioned: CASE expression routing to per-partition filters + /// + /// # Arguments + /// * `data` - Build data including hash map, pushdown strategy, and bounds + /// + /// # Returns + /// * `Result<()>` - Ok if successful, Err if filter update failed or mode mismatch + pub(crate) async fn report_build_data(&self, data: PartitionBuildData) -> Result<()> { + let finalize_input = { + let mut guard = self.inner.lock(); + self.store_build_data(&mut guard, data)?; + self.take_finalize_input_if_ready(&mut guard) + }; + + if let Some(finalize_input) = finalize_input { + self.finish(finalize_input); + } + + self.wait_for_completion().await + } + + pub(crate) fn report_canceled_partition(&self, partition_id: usize) { + let finalize_input = { + let mut guard = self.inner.lock(); + self.store_canceled_partition(&mut guard, partition_id); + self.take_finalize_input_if_ready(&mut guard) + }; + + if let Some(finalize_input) = finalize_input { + self.finish(finalize_input); + } + } + + fn store_build_data( + &self, + guard: &mut AccumulatorState, + data: PartitionBuildData, + ) -> Result<()> { + match (data, &mut guard.data) { + ( + PartitionBuildData::Partitioned { + partition_id, + pushdown, + bounds, + keys_have_null, + }, + AccumulatedBuildData::Partitioned { + partitions, + completed_partitions, + }, + ) => { + if matches!(partitions[partition_id], PartitionStatus::Pending) { + *completed_partitions += 1; + } + partitions[partition_id] = PartitionStatus::Reported(PartitionData { + pushdown, + bounds, + keys_have_null, + }); + } + ( + PartitionBuildData::CollectLeft { + pushdown, + bounds, + keys_have_null, + }, + AccumulatedBuildData::CollectLeft { + data, + reported_count, + .. + }, + ) => { + if matches!(data, PartitionStatus::Pending) { + *data = PartitionStatus::Reported(PartitionData { + pushdown, + bounds, + keys_have_null, + }); + } + *reported_count += 1; + } + _ => { + return datafusion_common::internal_err!( + "Build data mode mismatch in report_build_data" + ); + } + } + Ok(()) + } + + fn store_canceled_partition( + &self, + guard: &mut AccumulatorState, + partition_id: usize, + ) { + if let AccumulatedBuildData::Partitioned { + partitions, + completed_partitions, + } = &mut guard.data + && matches!(partitions[partition_id], PartitionStatus::Pending) + { + partitions[partition_id] = PartitionStatus::CanceledUnknown; + *completed_partitions += 1; + } + } + + fn take_finalize_input_if_ready( + &self, + guard: &mut AccumulatorState, + ) -> Option<FinalizeInput> { + if !matches!(guard.completion, CompletionState::Pending) { + return None; + } + + let finalize_input = match &guard.data { + AccumulatedBuildData::Partitioned { + partitions, + completed_partitions, + } if *completed_partitions == partitions.len() => { + Some(FinalizeInput::Partitioned(partitions.clone())) + } + AccumulatedBuildData::CollectLeft { + data, + reported_count, + expected_reports, + } if *reported_count == *expected_reports => { + Some(FinalizeInput::CollectLeft(data.clone())) + } + _ => None, + }?; + + guard.completion = CompletionState::Finalizing; + Some(finalize_input) + } + + fn finish(&self, finalize_input: FinalizeInput) { + let result = self.build_filter(finalize_input).map_err(Arc::new); + self.dynamic_filter.mark_complete(); + + let mut guard = self.inner.lock(); + guard.completion = CompletionState::Ready(result); + drop(guard); + self.completion_notify.notify_waiters(); + } + + async fn wait_for_completion(&self) -> Result<()> { + loop { + let notified = { + let guard = self.inner.lock(); + match &guard.completion { + CompletionState::Ready(Ok(())) => return Ok(()), + CompletionState::Ready(Err(err)) => { + return Err(DataFusionError::Shared(Arc::clone(err))); + } + CompletionState::Pending | CompletionState::Finalizing => { + self.completion_notify.notified() + } + } + }; + notified.await; + } + } + + fn build_filter(&self, finalize_input: FinalizeInput) -> Result<()> { + match finalize_input { + FinalizeInput::CollectLeft(partition) => match partition { + PartitionStatus::Reported(partition_data) => { + let membership_expr = create_membership_predicate( + &self.on_right, + partition_data.pushdown.clone(), + &HASH_JOIN_SEED, + self.probe_schema.as_ref(), + )?; + let bounds_expr = + create_bounds_predicate(&self.on_right, &partition_data.bounds); + + if let Some(filter_expr) = + combine_membership_and_bounds(membership_expr, bounds_expr) + { + self.dynamic_filter.update(self.preserve_probe_nulls( + filter_expr, + partition_data.keys_have_null, + )?)?; + } + } + PartitionStatus::Pending => { + return datafusion_common::internal_err!( + "attempted to finalize collect-left dynamic filter without reported build data" + ); + } + PartitionStatus::CanceledUnknown => { + return datafusion_common::internal_err!( + "collect-left dynamic filter cannot finalize with canceled build data" + ); + } + }, + FinalizeInput::Partitioned(partitions) => { + let num_partitions = partitions.len(); + let mut partition_filters = Vec::with_capacity(num_partitions); + let mut real_partition_ids = Vec::new(); + let mut empty_partition_ids = Vec::new(); + let mut has_canceled_unknown = false; + let mut keys_have_null = false; + + for (partition_id, partition) in partitions.iter().enumerate() { + match partition { + PartitionStatus::Reported(partition) + if matches!(partition.pushdown, PushdownStrategy::Empty) => + { + empty_partition_ids.push(partition_id); + partition_filters.push(lit(false)); + } + PartitionStatus::Reported(partition) => { + real_partition_ids.push(partition_id); + keys_have_null |= partition.keys_have_null; + let membership_expr = create_membership_predicate( + &self.on_right, + partition.pushdown.clone(), + &HASH_JOIN_SEED, + self.probe_schema.as_ref(), + )?; + let bounds_expr = create_bounds_predicate( + &self.on_right, + &partition.bounds, + ); + let then_expr = combine_membership_and_bounds( + membership_expr, + bounds_expr, + ) + .unwrap_or_else(|| lit(true)); + partition_filters.push(then_expr); + } + PartitionStatus::CanceledUnknown => { + has_canceled_unknown = true; + partition_filters.push(lit(true)); + // A canceled partition's build content is unknown, so it + // may hold a NULL key. + keys_have_null = true; + } + PartitionStatus::Pending => { + return datafusion_common::internal_err!( + "attempted to finalize dynamic filter with pending partition" + ); + } + } + } + + let filter_expr = if has_canceled_unknown + && real_partition_ids.is_empty() + && empty_partition_ids.is_empty() + { + lit(true) + } else if !has_canceled_unknown && real_partition_ids.is_empty() { + lit(false) + } else if !has_canceled_unknown + && real_partition_ids.len() == 1 + && empty_partition_ids.len() + 1 == num_partitions + { + Arc::clone(&partition_filters[real_partition_ids[0]]) + } else if let Some(range_partitioning) = &self.probe_range_partitioning { + // Range partitioning + assert_or_internal_err!( + partition_filters.len() == range_partitioning.partition_count(), + "Dynamic filter partition count {} does not match Range partition count {}", + partition_filters.len(), + range_partitioning.partition_count() + ); + let routing_range_expr = Arc::new(RangeExpr::try_new( + self.on_right.clone(), + range_partitioning, + )?) + as Arc<dyn PhysicalExpr>; + let else_expr = partition_filters + .pop() + .expect("Range partitioning always has at least one partition"); + + // CASE range_partition(key) + // WHEN 0 THEN F0 + // WHEN 1 THEN F1 + // ... + // ELSE Fn + // END + let when_then_expr = partition_filters + .into_iter() + .enumerate() + .map(|(partition_id, then_expr)| { + ( + lit(ScalarValue::UInt64(Some(partition_id as u64))), + then_expr, + ) + }) + .collect(); + + Arc::new(CaseExpr::try_new( + Some(routing_range_expr), + when_then_expr, + Some(else_expr), + )?) as Arc<dyn PhysicalExpr> + } else { + // Hash partitioning + let routing_hash_expr = Arc::new(HashExpr::new( + self.on_right.clone(), + self.repartition_random_state.clone(), + "hash_repartition".to_string(), + )) + as Arc<dyn PhysicalExpr>; + let modulo_expr = Arc::new(BinaryExpr::new( + routing_hash_expr, + Operator::Modulo, + lit(ScalarValue::UInt64(Some(num_partitions as u64))), + )) as Arc<dyn PhysicalExpr>; + + let mut when_then_branches = if has_canceled_unknown { + empty_partition_ids + .into_iter() + .map(|partition_id| { + ( + lit(ScalarValue::UInt64(Some(partition_id as u64))), + lit(false), + ) + }) + .collect::<Vec<_>>() + } else { + vec![] + }; + when_then_branches.extend(real_partition_ids.into_iter().map( + |partition_id| { + ( + lit(ScalarValue::UInt64(Some(partition_id as u64))), + Arc::clone(&partition_filters[partition_id]), + ) + }, + )); + + Arc::new(CaseExpr::try_new( + Some(modulo_expr), + when_then_branches, + Some(lit(has_canceled_unknown)), + )?) as Arc<dyn PhysicalExpr> + }; + + self.dynamic_filter + .update(self.preserve_probe_nulls(filter_expr, keys_have_null)?)?; + } + } + + Ok(()) + } + + /// Keeps probe rows with a NULL key when the join semantics need them. + /// + /// The build-side predicate drops probe rows whose key is NULL. A null-aware anti join + /// (`NOT IN`) needs that NULL to reach the join so three-valued logic can collapse the + /// result, and a null-equal join needs it to match a build-side NULL. OR-ing `key IS NULL` + /// keeps those rows while preserving the filter's selectivity for the rest; the join refines + /// whatever the widened filter lets through. + fn preserve_probe_nulls( + &self, + filter_expr: Arc<dyn PhysicalExpr>, + build_keys_have_null: bool, + ) -> Result<Arc<dyn PhysicalExpr>> { + // A null-aware anti join needs every probe NULL no matter what the build holds: one + // probe NULL makes `NOT IN` unknown for every build row. A null-equal join needs probe + // NULLs only to match an actual build-side NULL, so a NULL-free build keeps the filter + // at full selectivity. + let needs_probe_nulls = self.null_aware + || (self.null_equality == NullEquality::NullEqualsNull + && build_keys_have_null); + if !needs_probe_nulls { + return Ok(filter_expr); + } + // Only a key that can actually be NULL needs the disjunct; a NOT NULL key never widens. + // Null-aware joins are single-key; null-equal joins can be multi-key, so OR every nullable + // key. If every key is NOT NULL the filter is left untouched, at full selectivity. + let mut any_key_is_null: Option<Arc<dyn PhysicalExpr>> = None; + for key in &self.on_right { + // `nullable` fails only when a key is out of sync with the probe schema. That is + // a construction bug, so surface it instead of widening around it. + if !key.nullable(&self.probe_schema)? { + continue; + } + let is_null = + Arc::new(IsNullExpr::new(Arc::clone(key))) as Arc<dyn PhysicalExpr>; + any_key_is_null = Some(match any_key_is_null { + Some(acc) => Arc::new(BinaryExpr::new(acc, Operator::Or, is_null)) as _, + None => is_null, + }); + } + // Cheap null check first short-circuits before the costlier dynamic filter. + Ok(match any_key_is_null { + Some(any_key_is_null) => { + Arc::new(BinaryExpr::new(any_key_is_null, Operator::Or, filter_expr)) + } + None => filter_expr, + }) + } +} + +impl fmt::Debug for SharedBuildAccumulator { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SharedBuildAccumulator") + } +} + +#[cfg(test)] +pub(super) fn make_partitioned_accumulator_for_test( + num_partitions: usize, +) -> SharedBuildAccumulator { + let probe_schema = Arc::new(Schema::new(vec![Field::new( + "probe_key", + DataType::Int32, + false, + )])); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + SharedBuildAccumulator { + inner: Mutex::new(AccumulatorState { + data: AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; num_partitions], + completed_partitions: 0, + }, + completion: CompletionState::Pending, + }), + completion_notify: Notify::new(), + dynamic_filter, + on_right: vec![], + repartition_random_state: SeededRandomState::with_seed(1), + probe_schema, + probe_range_partitioning: None, + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + } +} + +#[cfg(test)] +pub(super) fn completed_partitions_for_test(acc: &SharedBuildAccumulator) -> usize { + let guard = acc.inner.lock(); + let AccumulatedBuildData::Partitioned { + completed_partitions, + .. + } = &guard.data + else { + panic!("expected partitioned accumulator"); + }; + *completed_partitions +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array}; + use arrow::compute::SortOptions; + use arrow::record_batch::RecordBatch; + use datafusion_common::SplitPoint; + use datafusion_physical_expr::{ + PhysicalSortExpr, + expressions::{Column, Literal}, + }; + + fn test_on_right() -> Vec<PhysicalExprRef> { + vec![Arc::new(Column::new("probe_key", 0))] + } + + fn test_probe_schema() -> Arc<Schema> { + Arc::new(Schema::new(vec![Field::new( + "probe_key", + DataType::Int32, + false, + )])) + } + + fn test_dynamic_filter( + on_right: &[PhysicalExprRef], + ) -> Arc<DynamicFilterPhysicalExpr> { + Arc::new(DynamicFilterPhysicalExpr::new(on_right.to_vec(), lit(true))) + } + + fn make_accumulator_for_test( + data: AccumulatedBuildData, + on_right: Vec<PhysicalExprRef>, + ) -> SharedBuildAccumulator { + let dynamic_filter = test_dynamic_filter(&on_right); + SharedBuildAccumulator { + inner: Mutex::new(AccumulatorState { + data, + completion: CompletionState::Pending, + }), + completion_notify: Notify::new(), + dynamic_filter, + on_right, + repartition_random_state: SeededRandomState::with_seed(1), + probe_schema: test_probe_schema(), + probe_range_partitioning: None, + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + } + } + + fn make_collect_left_accumulator_for_test() -> SharedBuildAccumulator { + make_accumulator_for_test( + AccumulatedBuildData::CollectLeft { + data: PartitionStatus::Pending, + reported_count: 0, + expected_reports: 1, + }, + test_on_right(), + ) + } + + fn make_partitioned_expr_accumulator_for_test( + num_partitions: usize, + ) -> SharedBuildAccumulator { + make_accumulator_for_test( + AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; num_partitions], + completed_partitions: 0, + }, + test_on_right(), + ) + } + + fn in_list(values: &[i32]) -> PushdownStrategy { + PushdownStrategy::InList(Arc::new(Int32Array::from(values.to_vec())) as ArrayRef) + } + + fn bounds(min: i32, max: i32) -> PartitionBounds { + PartitionBounds::new(vec![ColumnBounds::new( + ScalarValue::Int32(Some(min)), + ScalarValue::Int32(Some(max)), + )]) + } + + fn no_bounds() -> PartitionBounds { + PartitionBounds::new(vec![]) + } + + fn reported(pushdown: PushdownStrategy, bounds: PartitionBounds) -> PartitionStatus { + PartitionStatus::Reported(PartitionData { + pushdown, + bounds, + keys_have_null: false, + }) + } + + fn current_expr(acc: &SharedBuildAccumulator) -> PhysicalExprRef { + acc.dynamic_filter + .current() + .expect("dynamic filter current expression should be available") + } + + fn in_list_expr(expr: &PhysicalExprRef) -> &InListExpr { + expr.downcast_ref::<InListExpr>() + .expect("expected InListExpr dynamic filter") + } + + fn assert_in_list_column_values( + expr: &PhysicalExprRef, + expected_column_name: &str, + expected_column_index: usize, + expected_values: &[i32], + ) { + let in_list = in_list_expr(expr); + let column = in_list + .expr() + .downcast_ref::<Column>() + .expect("expected InListExpr child column"); + assert_eq!(column.name(), expected_column_name); + assert_eq!(column.index(), expected_column_index); + + let actual_values = in_list + .list() + .iter() + .map(|expr| { + let literal = expr + .downcast_ref::<Literal>() + .expect("expected InListExpr literal value"); + match literal.value() { + ScalarValue::Int32(Some(value)) => *value, + value => panic!("expected Int32 in-list value, got {value:?}"), + } + }) + .collect::<Vec<_>>(); + assert_eq!(actual_values, expected_values); + } + + fn binary_expr(expr: &PhysicalExprRef) -> &BinaryExpr { + expr.downcast_ref::<BinaryExpr>() + .expect("expected BinaryExpr dynamic filter") + } + + fn case_expr(expr: &PhysicalExprRef) -> &CaseExpr { + expr.downcast_ref::<CaseExpr>() + .expect("expected CaseExpr dynamic filter") + } + + fn assert_literal_bool(expr: &PhysicalExprRef, expected: bool) { + let literal = expr + .downcast_ref::<Literal>() + .expect("expected literal bool dynamic filter"); + assert_eq!(literal.value(), &ScalarValue::Boolean(Some(expected))); + } + + fn assert_top_binary_op(expr: &PhysicalExprRef, expected: Operator) { + assert_eq!(binary_expr(expr).op(), &expected); + } + + fn partitioned_state(acc: &SharedBuildAccumulator) -> (Vec<PartitionStatus>, usize) { + let guard = acc.inner.lock(); + let AccumulatedBuildData::Partitioned { + partitions, + completed_partitions, + } = &guard.data + else { + panic!("expected partitioned accumulator"); + }; + (partitions.clone(), *completed_partitions) + } + + #[test] + fn collect_left_updates_with_membership_only() { + let acc = make_collect_left_accumulator_for_test(); + + acc.build_filter(FinalizeInput::CollectLeft(reported( + in_list(&[1, 2, 3]), + no_bounds(), + ))) + .unwrap(); + + let expr = current_expr(&acc); + assert_in_list_column_values(&expr, "probe_key", 0, &[1, 2, 3]); + } + + #[test] + fn collect_left_updates_with_bounds_only() { + let acc = make_collect_left_accumulator_for_test(); + + acc.build_filter(FinalizeInput::CollectLeft(reported( + PushdownStrategy::Empty, + bounds(10, 20), + ))) + .unwrap(); + + let expr = current_expr(&acc); + assert_top_binary_op(&expr, Operator::And); + } + + #[test] + fn collect_left_empty_build_data_does_not_update_filter() { + let acc = make_collect_left_accumulator_for_test(); + let initial_generation = acc.dynamic_filter.snapshot_generation(); + + acc.build_filter(FinalizeInput::CollectLeft(reported( + PushdownStrategy::Empty, + no_bounds(), + ))) + .unwrap(); + + assert_eq!( + acc.dynamic_filter.snapshot_generation(), + initial_generation, + "empty CollectLeft input must not update with a no-op filter" + ); + let expr = current_expr(&acc); + assert_literal_bool(&expr, true); + } + + #[test] + fn partitioned_one_real_partition_with_rest_empty_skips_case() { + let acc = make_partitioned_expr_accumulator_for_test(3); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(PushdownStrategy::Empty, no_bounds()), + reported(in_list(&[2]), no_bounds()), + reported(PushdownStrategy::Empty, no_bounds()), + ])) + .unwrap(); + + let expr = current_expr(&acc); + in_list_expr(&expr); + assert!(expr.downcast_ref::<CaseExpr>().is_none()); + } + + #[test] + fn partitioned_canceled_unknown_partitions_keep_unknown_routes_permissive() { + let acc = make_partitioned_expr_accumulator_for_test(2); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + PartitionStatus::CanceledUnknown, + reported(PushdownStrategy::Empty, no_bounds()), + ])) + .unwrap(); + + let expr = current_expr(&acc); + let case = case_expr(&expr); + assert_eq!(case.when_then_expr().len(), 1); + assert_literal_bool(&case.when_then_expr()[0].1, false); + assert_literal_bool( + case.else_expr().expect("expected permissive fallback"), + true, + ); + } + + #[test] + fn partitioned_range_dynamic_filter_routes_with_range_expr() -> Result<()> { + let mut acc = make_partitioned_expr_accumulator_for_test(4); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + Default::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(PushdownStrategy::Empty, no_bounds()), + PartitionStatus::CanceledUnknown, + reported(in_list(&[20, 29]), no_bounds()), + reported(in_list(&[30]), no_bounds()), + ]))?; + + let expr = current_expr(&acc); + let case = case_expr(&expr); + assert!( + case.expr() + .and_then(|expr| expr.downcast_ref::<RangeExpr>()) + .is_some(), + "Range routing must use RangeExpr" + ); + assert_eq!(case.when_then_expr().len(), 3); + + let batch = RecordBatch::try_new( + test_probe_schema(), + vec![Arc::new(Int32Array::from(vec![ + 9, 10, 19, 20, 21, 29, 30, 31, + ]))], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = result + .as_any() + .downcast_ref::<BooleanArray>() + .expect("dynamic filter should evaluate to BooleanArray"); + assert_eq!( + result, + &BooleanArray::from(vec![false, true, true, true, false, true, true, false,]) + ); + + Ok(()) + } + + #[test] + fn partitioned_range_dynamic_filter_routes_compound_nullable_keys() -> Result<()> { + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("probe_key", DataType::Int32, true), + Field::new("probe_tie", DataType::Int32, true), + ])); + let on_right: Vec<PhysicalExprRef> = vec![ + Arc::new(Column::new("probe_key", 0)), + Arc::new(Column::new("probe_tie", 1)), + ]; + let mut acc = make_accumulator_for_test( + AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; 4], + completed_partitions: 0, + }, + on_right, + ); + acc.probe_schema = Arc::clone(&probe_schema); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [ + PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::clone(&acc.on_right[1]), + SortOptions::new(false, false), + ), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(None), + ScalarValue::Int32(Some(10)), + ]), + SplitPoint::new(vec![ScalarValue::Int32(None), ScalarValue::Int32(None)]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(None), + ]), + ], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(PushdownStrategy::Empty, no_bounds()), + PartitionStatus::CanceledUnknown, + reported(PushdownStrategy::Empty, no_bounds()), + PartitionStatus::CanceledUnknown, + ]))?; + + let expr = current_expr(&acc); + let case = case_expr(&expr); + assert!(case.expr().is_some()); + assert_eq!(case.when_then_expr().len(), 3); + + let batch = RecordBatch::try_new( + probe_schema, + vec![ + Arc::new(Int32Array::from(vec![ + None, + None, + None, + None, + Some(9), + Some(10), + Some(10), + Some(11), + ])), + Arc::new(Int32Array::from(vec![ + Some(9), + Some(10), + Some(11), + None, + None, + Some(9), + None, + None, + ])), + ], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = result + .as_any() + .downcast_ref::<BooleanArray>() + .expect("dynamic filter should evaluate to BooleanArray"); + assert_eq!( + result, + &BooleanArray::from( + vec![false, true, true, false, false, false, true, true,] + ) + ); + + Ok(()) + } + + #[test] + fn partitioned_range_dynamic_filter_preserves_signed_zero_routing() -> Result<()> { + let probe_schema = Arc::new(Schema::new(vec![Field::new( + "probe_key", + DataType::Float64, + false, + )])); + let on_right: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("probe_key", 0))]; + let mut acc = make_accumulator_for_test( + AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; 2], + completed_partitions: 0, + }, + on_right, + ); + acc.probe_schema = Arc::clone(&probe_schema); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + SortOptions::default(), + )] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Float64(Some(0.0))])], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + PartitionStatus::CanceledUnknown, + reported(PushdownStrategy::Empty, no_bounds()), + ]))?; + + let expr = current_expr(&acc); + let batch = RecordBatch::try_new( + probe_schema, + vec![Arc::new(Float64Array::from(vec![-0.0, 0.0]))], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = result + .as_any() + .downcast_ref::<BooleanArray>() + .expect("dynamic filter should evaluate to BooleanArray"); + assert_eq!(result, &BooleanArray::from(vec![true, false])); + + Ok(()) + } + + // Regression guard for the build-report lifecycle fix: on `Drop`, a stream + // in `BuildReportState::ReportScheduled` still calls `report_canceled_partition` + // because it cannot tell whether the coordinator has already observed the + // report (first poll of the `OnceFut` runs `store_build_data` synchronously + // before the future's first `.await`, but the stream doesn't learn that + // until `get_shared` returns `Ok`). Correctness therefore relies on + // `store_canceled_partition` being a no-op when the partition is already + // `Reported`. This test pins that invariant. + #[test] + fn report_canceled_partition_is_noop_after_report() { + let acc = make_partitioned_accumulator_for_test(2); + + { + let mut guard = acc.inner.lock(); + acc.store_build_data( + &mut guard, + PartitionBuildData::Partitioned { + partition_id: 0, + pushdown: PushdownStrategy::Empty, + bounds: PartitionBounds::new(vec![]), + keys_have_null: false, + }, + ) + .unwrap(); + } + let (partitions, completed) = partitioned_state(&acc); + assert!(matches!(partitions[0], PartitionStatus::Reported(_))); + assert_eq!(completed, 1); + + acc.report_canceled_partition(0); + let (partitions, completed) = partitioned_state(&acc); + assert!( + matches!(partitions[0], PartitionStatus::Reported(_)), + "late cancel must not overwrite a prior Reported status" + ); + assert_eq!(completed, 1, "late cancel must not double-count completion"); + } + + // Drop from the `NotReported` (or first-poll-never-ran) state must + // transition `Pending` -> `CanceledUnknown` and bump `completed_partitions`, + // which is what unblocks sibling partitions waiting on the coordinator. + #[test] + fn report_canceled_partition_marks_pending_partition_canceled() { + let acc = make_partitioned_accumulator_for_test(2); + + acc.report_canceled_partition(0); + let (partitions, completed) = partitioned_state(&acc); + assert!(matches!(partitions[0], PartitionStatus::CanceledUnknown)); + assert_eq!(completed, 1); + + // Idempotent: a second cancel (e.g. a stray double-drop) must not + // double-count completion. + acc.report_canceled_partition(0); + let (partitions, completed) = partitioned_state(&acc); + assert!(matches!(partitions[0], PartitionStatus::CanceledUnknown)); + assert_eq!(completed, 1); + } + + fn null_semantics_accumulator( + probe_schema: Arc<Schema>, + on_right: Vec<PhysicalExprRef>, + null_equality: NullEquality, + null_aware: bool, + ) -> SharedBuildAccumulator { + SharedBuildAccumulator { + inner: Mutex::new(AccumulatorState { + data: AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; 1], + completed_partitions: 0, + }, + completion: CompletionState::Pending, + }), + completion_notify: Notify::new(), + dynamic_filter: Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))), + on_right, + repartition_random_state: SeededRandomState::with_seed(1), + probe_schema, + probe_range_partitioning: None, + null_equality, + null_aware, + } + } + + fn null_equal_accumulator( + probe_schema: Arc<Schema>, + on_right: Vec<PhysicalExprRef>, + ) -> SharedBuildAccumulator { + null_semantics_accumulator( + probe_schema, + on_right, + NullEquality::NullEqualsNull, + false, + ) + } + + #[test] + fn preserve_probe_nulls_only_widens_nullable_keys() { + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("k_nullable", DataType::Int32, true), + Field::new("k_not_null", DataType::Int32, false), + ])); + let on_right: Vec<PhysicalExprRef> = vec![ + Arc::new(Column::new("k_nullable", 0)), + Arc::new(Column::new("k_not_null", 1)), + ]; + let acc = null_equal_accumulator(probe_schema, on_right); + + // Only the nullable key earns an IS NULL disjunct; the NOT NULL key is left out. + let widened = acc.preserve_probe_nulls(lit(true), true).unwrap(); + assert_eq!(format!("{widened}").matches("IS NULL").count(), 1); + } + + #[test] + fn preserve_probe_nulls_leaves_all_not_null_keys_untouched() { + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let on_right: Vec<PhysicalExprRef> = + vec![Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1))]; + let acc = null_equal_accumulator(probe_schema, on_right); + + // Every key is NOT NULL, so there is nothing to OR in and the filter is returned as-is. + let filter = lit(true); + let result = acc.preserve_probe_nulls(Arc::clone(&filter), true).unwrap(); + assert_eq!(format!("{result}"), format!("{filter}")); + } + + #[test] + fn preserve_probe_nulls_rejects_out_of_sync_key() { + let probe_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + // The key's column index points past the probe schema: a construction bug that + // must surface as an error, not get widened around. + let on_right: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b", 1))]; + let acc = null_equal_accumulator(probe_schema, on_right); + + assert!(acc.preserve_probe_nulls(lit(true), true).is_err()); + } + + #[test] + fn preserve_probe_nulls_skips_wrap_when_build_has_no_nulls() { + let probe_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let on_right: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("a", 0))]; + let acc = null_equal_accumulator(probe_schema, on_right); + + // A NULL-free build has nothing for a probe NULL to null-match, so the + // filter keeps its full selectivity. + let filter = lit(true); + let result = acc + .preserve_probe_nulls(Arc::clone(&filter), false) + .unwrap(); + assert_eq!(format!("{result}"), format!("{filter}")); + } + + #[test] + fn preserve_probe_nulls_wraps_null_aware_regardless_of_build() { + let probe_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let on_right: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("a", 0))]; + let acc = null_semantics_accumulator( + probe_schema, + on_right, + NullEquality::NullEqualsNothing, + true, + ); + + // One probe NULL collapses `NOT IN` for every build row, so the wrap must not + // depend on the build content. + let widened = acc.preserve_probe_nulls(lit(true), false).unwrap(); + assert_eq!(format!("{widened}").matches("IS NULL").count(), 1); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs b/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs new file mode 100644 index 0000000..6869395 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/hash_join/stream.rs
@@ -0,0 +1,1144 @@ +// 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. + +//! Stream implementation for Hash Join +//! +//! This module implements [`HashJoinStream`], the streaming engine for +//! [`super::HashJoinExec`]. See comments in [`HashJoinStream`] for more details. + +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::task::Poll; + +use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; +use crate::joins::Map; +use crate::joins::MapOffset; +use crate::joins::PartitionMode; +use crate::joins::hash_join::exec::JoinLeftData; +use crate::joins::hash_join::shared_bounds::{ + PartitionBounds, PartitionBuildData, SharedBuildAccumulator, +}; +use crate::joins::utils::{ + OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, matchable_join_keys, +}; +use crate::stream::EmptyRecordBatchStream; +use crate::{ + RecordBatchStream, SendableRecordBatchStream, handle_state, + hash_utils::create_hashes, + joins::utils::{ + BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, + StatefulStreamResult, adjust_indices_by_join_type, apply_join_filter_to_indices, + build_batch_empty_build_side, build_batch_from_indices, + need_produce_result_in_final, + }, +}; + +use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array}; +use arrow::buffer::NullBuffer; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{ + JoinSide, JoinType, NullEquality, Result, internal_datafusion_err, internal_err, +}; +use datafusion_physical_expr::PhysicalExprRef; + +use datafusion_common::hash_utils::RandomState; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::{Stream, StreamExt, ready}; + +/// Represents build-side of hash join. +pub(super) enum BuildSide { + /// Indicates that build-side not collected yet + Initial(BuildSideInitialState), + /// Indicates that build-side data has been collected + Ready(BuildSideReadyState), +} + +/// Container for BuildSide::Initial related data +pub(super) struct BuildSideInitialState { + /// Future for building hash table from build-side input + pub(super) left_fut: OnceFut<JoinLeftData>, +} + +/// Container for BuildSide::Ready related data +pub(super) struct BuildSideReadyState { + /// Collected build-side data + left_data: Arc<JoinLeftData>, +} + +impl BuildSide { + /// Tries to extract BuildSideInitialState from BuildSide enum. + /// Returns an error if state is not Initial. + fn try_as_initial_mut(&mut self) -> Result<&mut BuildSideInitialState> { + match self { + BuildSide::Initial(state) => Ok(state), + _ => internal_err!("Expected build side in initial state"), + } + } + + /// Tries to extract BuildSideReadyState from BuildSide enum. + /// Returns an error if state is not Ready. + fn try_as_ready(&self) -> Result<&BuildSideReadyState> { + match self { + BuildSide::Ready(state) => Ok(state), + _ => internal_err!("Expected build side in ready state"), + } + } + + /// Tries to extract BuildSideReadyState from BuildSide enum. + /// Returns an error if state is not Ready. + fn try_as_ready_mut(&mut self) -> Result<&mut BuildSideReadyState> { + match self { + BuildSide::Ready(state) => Ok(state), + _ => internal_err!("Expected build side in ready state"), + } + } +} + +/// Represents state of HashJoinStream +/// +/// Expected state transitions performed by HashJoinStream are: +/// +/// ```text +/// +/// WaitBuildSide +/// │ +/// ▼ +/// ┌─► FetchProbeBatch ───► ExhaustedProbeSide ───► Completed +/// │ │ +/// │ ▼ +/// └─ ProcessProbeBatch +/// ``` +#[derive(Debug, Clone)] +pub(super) enum HashJoinStreamState { + /// Initial state for HashJoinStream indicating that build-side data not collected yet + WaitBuildSide, + /// Waiting for bounds to be reported by all partitions + WaitPartitionBoundsReport, + /// Indicates that build-side has been collected, and stream is ready for fetching probe-side + FetchProbeBatch, + /// Indicates that non-empty batch has been fetched from probe-side, and is ready to be processed + ProcessProbeBatch(ProcessProbeBatchState), + /// Indicates that probe-side has been fully processed + ExhaustedProbeSide, + /// Indicates that HashJoinStream execution is completed + Completed, +} + +impl HashJoinStreamState { + /// Tries to extract ProcessProbeBatchState from HashJoinStreamState enum. + /// Returns an error if state is not ProcessProbeBatchState. + fn try_as_process_probe_batch_mut(&mut self) -> Result<&mut ProcessProbeBatchState> { + match self { + HashJoinStreamState::ProcessProbeBatch(state) => Ok(state), + _ => internal_err!("Expected hash join stream in ProcessProbeBatch state"), + } + } +} + +/// Container for HashJoinStreamState::ProcessProbeBatch related data +#[derive(Debug, Clone)] +pub(super) struct ProcessProbeBatchState { + /// Current probe-side batch + batch: RecordBatch, + /// Probe-side on expressions values + values: Vec<ArrayRef>, + /// Combined validity of the probe-side key columns, set when NULL keys + /// exist and cannot match (`NullEquality::NullEqualsNothing`); NULL rows + /// are skipped during JoinHashMap lookups + valid_keys: Option<NullBuffer>, + /// Starting offset for JoinHashMap lookups + offset: MapOffset, + /// Max joined probe-side index from current batch + joined_probe_idx: Option<usize>, +} + +impl ProcessProbeBatchState { + fn advance(&mut self, offset: MapOffset, joined_probe_idx: Option<usize>) { + self.offset = offset; + if joined_probe_idx.is_some() { + self.joined_probe_idx = joined_probe_idx; + } + } +} + +/// Lifecycle of this partition's build-data report to the shared coordinator. +/// +/// `Scheduled` means the reporting `OnceFut` has been constructed but is lazy: +/// the coordinator has not necessarily observed the report. Only `Delivered` +/// guarantees the coordinator saw it, so `Drop` must still cancel a `Scheduled` +/// partition — otherwise sibling partitions can wait forever for a report that +/// never runs. +#[derive(Debug, PartialEq, Eq)] +enum BuildReportState { + NotReported, + Scheduled, + Delivered, + Canceled, + Finalized, +} + +/// Owns the stream-side lifecycle for one partition's build-data report. +struct BuildReportHandle { + partition: usize, + mode: PartitionMode, + build_accumulator: Option<Arc<SharedBuildAccumulator>>, + waiter: Option<OnceFut<()>>, + state: BuildReportState, +} + +impl BuildReportHandle { + fn new( + partition: usize, + mode: PartitionMode, + build_accumulator: Option<Arc<SharedBuildAccumulator>>, + ) -> Self { + Self { + partition, + mode, + build_accumulator, + waiter: None, + state: BuildReportState::NotReported, + } + } + + fn has_accumulator(&self) -> bool { + self.build_accumulator.is_some() + } + + fn schedule(&mut self, build_data: PartitionBuildData) { + let Some(build_accumulator) = &self.build_accumulator else { + // Defensive no-op terminal state; current callers avoid scheduling + // unless an accumulator is present. + self.finalize(); + return; + }; + + debug_assert!(matches!(self.state, BuildReportState::NotReported)); + let acc = Arc::clone(build_accumulator); + self.waiter = Some(OnceFut::new(async move { + acc.report_build_data(build_data).await + })); + self.state = BuildReportState::Scheduled; + } + + fn poll_delivery(&mut self, cx: &mut std::task::Context<'_>) -> Poll<Result<()>> { + if let Some(ref mut fut) = self.waiter { + ready!(fut.get_shared(cx))?; + if !matches!(self.state, BuildReportState::Delivered) { + debug_assert!(matches!(self.state, BuildReportState::Scheduled)); + self.state = BuildReportState::Delivered; + } + } + Poll::Ready(Ok(())) + } + + fn cancel_pending(&mut self) { + if matches!( + self.state, + BuildReportState::Delivered + | BuildReportState::Canceled + | BuildReportState::Finalized + ) { + return; + } + + if self.mode == PartitionMode::Partitioned + && let Some(build_accumulator) = &self.build_accumulator + { + build_accumulator.report_canceled_partition(self.partition); + self.state = BuildReportState::Canceled; + } else { + self.finalize(); + } + } + + fn finalize(&mut self) { + self.state = BuildReportState::Finalized; + } + + #[cfg(test)] + fn state(&self) -> &BuildReportState { + &self.state + } +} + +impl Drop for BuildReportHandle { + fn drop(&mut self) { + self.cancel_pending(); + } +} + +/// [`Stream`] for [`super::HashJoinExec`] that does the actual join. +/// +/// This stream: +/// +/// - Collecting the build side (left input) into a hash map +/// - Iterating over the probe side (right input) in streaming fashion +/// - Looking up matches against the hash table and applying join filters +/// - Producing joined [`RecordBatch`]es incrementally +/// - Emitting unmatched rows for outer/semi/anti joins in the final stage +pub(super) struct HashJoinStream { + /// Partition identifier for debugging and determinism + partition: usize, + /// Input schema + schema: Arc<Schema>, + /// equijoin columns from the right (probe side) + on_right: Vec<PhysicalExprRef>, + /// optional join filter + filter: Option<JoinFilter>, + /// type of the join (left, right, semi, etc) + join_type: JoinType, + /// right (probe) input + right: SendableRecordBatchStream, + /// Random state used for hashing initialization + random_state: RandomState, + /// Metrics + join_metrics: BuildProbeJoinMetrics, + /// Information of index and left / right placement of columns + column_indices: Vec<ColumnIndex>, + /// Defines the null equality for the join. + null_equality: NullEquality, + /// State of the stream + state: HashJoinStreamState, + /// Build side + build_side: BuildSide, + /// Maximum output batch size + batch_size: usize, + /// Scratch space for computing hashes + hashes_buffer: Vec<u64>, + /// Scratch space for probe indices during hash lookup + probe_indices_buffer: Vec<u32>, + /// Scratch space for build indices during hash lookup + build_indices_buffer: Vec<u64>, + /// Specifies whether the right side has an ordering to potentially preserve + right_side_ordered: bool, + /// Owns this partition's build-data report lifecycle. + build_report: BuildReportHandle, + /// Partitioning mode to use + mode: PartitionMode, + /// Output buffer for coalescing small batches into larger ones with optional fetch limit. + /// Uses `LimitedBatchCoalescer` to efficiently combine batches and absorb limit with 'fetch' + output_buffer: LimitedBatchCoalescer, + /// Whether this is a null-aware anti join + null_aware: bool, +} + +impl RecordBatchStream for HashJoinStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Executes lookups by hash against JoinHashMap and resolves potential +/// hash collisions. +/// Returns build/probe indices satisfying the equality condition, along with +/// (optional) starting point for next iteration. +/// +/// # Example +/// +/// For `LEFT.b1 = RIGHT.b2`: +/// LEFT (build) Table: +/// ```text +/// a1 b1 c1 +/// 1 1 10 +/// 3 3 30 +/// 5 5 50 +/// 7 7 70 +/// 9 8 90 +/// 11 8 110 +/// 13 10 130 +/// ``` +/// +/// RIGHT (probe) Table: +/// ```text +/// a2 b2 c2 +/// 2 2 20 +/// 4 4 40 +/// 6 6 60 +/// 8 8 80 +/// 10 10 100 +/// 12 10 120 +/// ``` +/// +/// The result is +/// ```text +/// "+----+----+-----+----+----+-----+", +/// "| a1 | b1 | c1 | a2 | b2 | c2 |", +/// "+----+----+-----+----+----+-----+", +/// "| 9 | 8 | 90 | 8 | 8 | 80 |", +/// "| 11 | 8 | 110 | 8 | 8 | 80 |", +/// "| 13 | 10 | 130 | 10 | 10 | 100 |", +/// "| 13 | 10 | 130 | 12 | 10 | 120 |", +/// "+----+----+-----+----+----+-----+" +/// ``` +/// +/// And the result of build and probe indices are: +/// ```text +/// Build indices: 4, 5, 6, 6 +/// Probe indices: 3, 3, 4, 5 +/// ``` +#[expect(clippy::too_many_arguments)] +pub(super) fn lookup_join_hashmap( + build_hashmap: &dyn JoinHashMapType, + build_side_values: &[ArrayRef], + probe_side_values: &[ArrayRef], + null_equality: NullEquality, + hashes_buffer: &[u64], + valid_keys: Option<&NullBuffer>, + limit: usize, + offset: MapOffset, + probe_indices_buffer: &mut Vec<u32>, + build_indices_buffer: &mut Vec<u64>, +) -> Result<(UInt64Array, UInt32Array, Option<MapOffset>)> { + let next_offset = build_hashmap.get_matched_indices_with_limit_offset( + hashes_buffer, + valid_keys, + limit, + offset, + probe_indices_buffer, + build_indices_buffer, + ); + + let build_indices_unfiltered: UInt64Array = + std::mem::take(build_indices_buffer).into(); + let probe_indices_unfiltered: UInt32Array = + std::mem::take(probe_indices_buffer).into(); + + // TODO: optimize equal_rows_arr to avoid allocation of intermediate arrays + // https://github.com/apache/datafusion/issues/12131 + let (build_indices, probe_indices) = equal_rows_arr( + &build_indices_unfiltered, + &probe_indices_unfiltered, + build_side_values, + probe_side_values, + null_equality, + )?; + + // Reclaim buffers + *build_indices_buffer = build_indices_unfiltered.into_parts().1.into(); + *probe_indices_buffer = probe_indices_unfiltered.into_parts().1.into(); + + Ok((build_indices, probe_indices, next_offset)) +} + +/// Counts the number of distinct elements in the input array. +/// +/// The input array must be sorted (e.g., `[0, 1, 1, 2, 2, ...]`) and contain no null values. +#[inline] +fn count_distinct_sorted_indices(indices: &UInt32Array) -> usize { + if indices.is_empty() { + return 0; + } + + debug_assert!(indices.null_count() == 0); + + let values_buf = indices.values(); + let values = values_buf.as_ref(); + let mut iter = values.iter(); + let Some(&first) = iter.next() else { + return 0; + }; + + let mut count = 1usize; + let mut last = first; + for &value in iter { + if value != last { + last = value; + count += 1; + } + } + count +} + +impl HashJoinStream { + #[expect(clippy::too_many_arguments)] + pub(super) fn new( + partition: usize, + schema: Arc<Schema>, + on_right: Vec<PhysicalExprRef>, + filter: Option<JoinFilter>, + join_type: JoinType, + right: SendableRecordBatchStream, + random_state: RandomState, + join_metrics: BuildProbeJoinMetrics, + column_indices: Vec<ColumnIndex>, + null_equality: NullEquality, + state: HashJoinStreamState, + build_side: BuildSide, + batch_size: usize, + hashes_buffer: Vec<u64>, + right_side_ordered: bool, + build_accumulator: Option<Arc<SharedBuildAccumulator>>, + mode: PartitionMode, + null_aware: bool, + fetch: Option<usize>, + ) -> Self { + // Create output buffer with coalescing and optional fetch limit. + let output_buffer = + LimitedBatchCoalescer::new(Arc::clone(&schema), batch_size, fetch); + + Self { + partition, + schema, + on_right, + filter, + join_type, + right, + random_state, + join_metrics, + column_indices, + null_equality, + state, + build_side, + batch_size, + hashes_buffer, + probe_indices_buffer: Vec::with_capacity(batch_size), + build_indices_buffer: Vec::with_capacity(batch_size), + right_side_ordered, + build_report: BuildReportHandle::new(partition, mode, build_accumulator), + mode, + output_buffer, + null_aware, + } + } + + /// Returns the next state after the build side has been fully collected + /// and any required build-side coordination has completed. + fn state_after_build_ready( + join_type: JoinType, + left_data: &JoinLeftData, + ) -> HashJoinStreamState { + let build_empty = !left_data.has_build_rows(); + // The map can be empty even when the build side has rows: under + // `NullEqualsNothing`, build rows with a NULL join key are omitted. For + // join types whose every output row requires a build match, that still + // guarantees an empty result, so we can skip scanning the probe side. + let map_empty = !left_data.has_matchable_build_rows(); + + if (build_empty && join_type.empty_build_side_produces_empty_result()) + || (map_empty && join_type.empty_map_produces_empty_result()) + { + HashJoinStreamState::Completed + } else { + HashJoinStreamState::FetchProbeBatch + } + } + + /// Transitions state after build-side data has been collected, automatically + /// reporting build data to the accumulator when one is present. + /// + /// If a `build_accumulator` is configured, this method constructs the + /// appropriate [`PartitionBuildData`], schedules the reporting future, and + /// returns [`HashJoinStreamState::WaitPartitionBoundsReport`]. Otherwise it + /// delegates to [`Self::state_after_build_ready`]. + fn transition_after_build_collected( + &mut self, + left_data: &Arc<JoinLeftData>, + ) -> HashJoinStreamState { + if !self.build_report.has_accumulator() { + return Self::state_after_build_ready(self.join_type, left_data.as_ref()); + } + + let pushdown = left_data.membership().clone(); + let bounds = left_data + .bounds + .clone() + .unwrap_or_else(|| PartitionBounds::new(vec![])); + // Arrow tracks null counts per array, so this costs no data scan. + let keys_have_null = left_data + .values() + .iter() + .any(|array| array.null_count() > 0); + + let build_data = match self.mode { + PartitionMode::Partitioned => PartitionBuildData::Partitioned { + partition_id: self.partition, + pushdown, + bounds, + keys_have_null, + }, + PartitionMode::CollectLeft => PartitionBuildData::CollectLeft { + pushdown, + bounds, + keys_have_null, + }, + PartitionMode::Auto => unreachable!( + "PartitionMode::Auto should not be present at execution time. This is a bug in DataFusion, please report it!" + ), + }; + + self.build_report.schedule(build_data); + HashJoinStreamState::WaitPartitionBoundsReport + } + + /// Separate implementation function that unpins the [`HashJoinStream`] so + /// that partial borrows work correctly + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + // First, check if we have any completed batches ready to emit + if let Some(batch) = self.output_buffer.next_completed_batch() { + return self + .join_metrics + .baseline + .record_poll(Poll::Ready(Some(Ok(batch)))); + } + + // Check if the coalescer has finished (limit reached and flushed) + if self.output_buffer.is_finished() { + return Poll::Ready(None); + } + + return match self.state { + HashJoinStreamState::WaitBuildSide => { + handle_state!(ready!(self.collect_build_side(cx))) + } + HashJoinStreamState::WaitPartitionBoundsReport => { + handle_state!(ready!(self.wait_for_partition_bounds_report(cx))) + } + HashJoinStreamState::FetchProbeBatch => { + handle_state!(ready!(self.fetch_probe_batch(cx))) + } + HashJoinStreamState::ProcessProbeBatch(_) => { + handle_state!(self.process_probe_batch()) + } + HashJoinStreamState::ExhaustedProbeSide => { + handle_state!(self.process_unmatched_build_batch()) + } + HashJoinStreamState::Completed if !self.output_buffer.is_empty() => { + // Flush any remaining buffered data + self.output_buffer.finish()?; + // Continue loop to emit the flushed batch + continue; + } + HashJoinStreamState::Completed => Poll::Ready(None), + }; + } + } + + /// Optional step to wait until build-side information (hash maps or bounds) has been reported by all partitions. + /// This state is only entered if a build accumulator is present. + /// + /// ## Why wait? + /// + /// The dynamic filter is only built once all partitions have reported their information (hash maps or bounds). + /// If we do not wait here, the probe-side scan may start before the filter is ready. + /// This can lead to the probe-side scan missing the opportunity to apply the filter + /// and skip reading unnecessary data. + fn wait_for_partition_bounds_report( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + ready!(self.build_report.poll_delivery(cx))?; + let build_side = self.build_side.try_as_ready()?; + self.state = + Self::state_after_build_ready(self.join_type, build_side.left_data.as_ref()); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Collects build-side data by polling `OnceFut` future from initialized build-side + /// + /// Updates build-side to `Ready`, and state to `FetchProbeSide` + fn collect_build_side( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + let build_timer = self.join_metrics.build_time.timer(); + // build hash table from left (build) side, if not yet done + let left_data = ready!( + self.build_side + .try_as_initial_mut()? + .left_fut + .get_shared(cx) + )?; + build_timer.done(); + + // Note: For null-aware anti join, we need to check the probe side (right) for NULLs, + // not the build side (left). The probe-side NULL check happens during process_probe_batch. + // The probe_side_has_null flag will be set there if any probe batch contains NULL. + + self.state = self.transition_after_build_collected(&left_data); + + self.build_side = BuildSide::Ready(BuildSideReadyState { left_data }); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Fetches next batch from probe-side + /// + /// If non-empty batch has been fetched, updates state to `ProcessProbeBatchState`, + /// otherwise updates state to `ExhaustedProbeSide` + fn fetch_probe_batch( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.right.poll_next_unpin(cx)) { + None => { + // Release the probe-side input pipeline's resources. The schema + // is preserved so callers that still query `self.right.schema()` + // (e.g. for unmatched-build emission) keep working. + let right_schema = self.right.schema(); + self.right = Box::pin(EmptyRecordBatchStream::new(right_schema)); + self.state = HashJoinStreamState::ExhaustedProbeSide; + } + Some(Ok(batch)) => { + // Precalculate hash values for fetched batch + let keys_values = evaluate_expressions_to_arrays(&self.on_right, &batch)?; + + let valid_keys = if let Map::HashMap(_) = + self.build_side.try_as_ready()?.left_data.map() + { + self.hashes_buffer.clear(); + self.hashes_buffer.resize(batch.num_rows(), 0); + create_hashes( + &keys_values, + &self.random_state, + &mut self.hashes_buffer, + )?; + matchable_join_keys(&keys_values, self.null_equality) + } else { + None + }; + + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(batch.num_rows()); + + self.state = + HashJoinStreamState::ProcessProbeBatch(ProcessProbeBatchState { + batch, + values: keys_values, + valid_keys, + offset: (0, None), + joined_probe_idx: None, + }); + } + Some(Err(err)) => return Poll::Ready(Err(err)), + }; + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + /// Joins current probe batch with build-side data and produces batch with matched output + /// + /// Updates state to `FetchProbeBatch` + fn process_probe_batch( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + let state = self.state.try_as_process_probe_batch_mut()?; + let build_side = self.build_side.try_as_ready_mut()?; + + self.join_metrics + .probe_hit_rate + .add_total(state.batch.num_rows()); + + let timer = self.join_metrics.join_time.timer(); + + // Null-aware anti join semantics: + // For LeftAnti: output LEFT (build) rows where LEFT.key NOT IN RIGHT.key + // 1. If RIGHT (probe) contains NULL in any batch, no LEFT rows should be output + // 2. LEFT rows with NULL keys should not be output (handled in final stage) + if self.null_aware { + // Mark that we've seen a probe batch with actual rows (probe side is non-empty) + // Only set this if batch has rows - empty batches don't count + // Use shared atomic state so all partitions can see this global information + if state.batch.num_rows() > 0 { + build_side + .left_data + .probe_side_non_empty + .store(true, Ordering::Relaxed); + } + + // Check if probe side (RIGHT) contains NULL + // Since null_aware validation ensures single column join, we only check the first column + let probe_key_column = &state.values[0]; + if probe_key_column.null_count() > 0 { + // Found NULL in probe side - set shared flag to prevent any output + build_side + .left_data + .probe_side_has_null + .store(true, Ordering::Relaxed); + } + + // If probe side has NULL (detected in this or any other partition), return empty result + if build_side + .left_data + .probe_side_has_null + .load(Ordering::Relaxed) + { + timer.done(); + self.state = HashJoinStreamState::FetchProbeBatch; + return Ok(StatefulStreamResult::Continue); + } + } + + let is_empty = !build_side.left_data.has_matchable_build_rows(); + + if is_empty { + let result = build_batch_empty_build_side( + &self.schema, + build_side.left_data.batch(), + &state.batch, + &self.column_indices, + self.join_type, + )?; + timer.done(); + self.output_buffer.push_batch(result)?; + self.state = HashJoinStreamState::FetchProbeBatch; + + return Ok(StatefulStreamResult::Continue); + } + + // get the matched by join keys indices + let (left_indices, right_indices, next_offset) = match build_side.left_data.map() + { + Map::HashMap(map) => lookup_join_hashmap( + map.as_ref(), + build_side.left_data.values(), + &state.values, + self.null_equality, + &self.hashes_buffer, + state.valid_keys.as_ref(), + self.batch_size, + state.offset, + &mut self.probe_indices_buffer, + &mut self.build_indices_buffer, + )?, + Map::ArrayMap(array_map) => { + let next_offset = array_map.get_matched_indices_with_limit_offset( + &state.values, + self.batch_size, + state.offset, + &mut self.probe_indices_buffer, + &mut self.build_indices_buffer, + )?; + ( + UInt64Array::from(self.build_indices_buffer.clone()), + UInt32Array::from(self.probe_indices_buffer.clone()), + next_offset, + ) + } + }; + + let distinct_right_indices_count = count_distinct_sorted_indices(&right_indices); + + self.join_metrics + .probe_hit_rate + .add_part(distinct_right_indices_count); + + self.join_metrics.avg_fanout.add_part(left_indices.len()); + + self.join_metrics + .avg_fanout + .add_total(distinct_right_indices_count); + + // apply join filter if exists + let (left_indices, right_indices) = if let Some(filter) = &self.filter { + apply_join_filter_to_indices( + build_side.left_data.batch(), + &state.batch, + left_indices, + right_indices, + filter, + JoinSide::Left, + None, + self.join_type, + )? + } else { + (left_indices, right_indices) + }; + + // mark joined left-side indices as visited, if required by join type + if need_produce_result_in_final(self.join_type) { + let mut bitmap = build_side.left_data.visited_indices_bitmap().lock(); + left_indices.iter().flatten().for_each(|x| { + bitmap.set_bit(x as usize, true); + }); + } + + // The goals of index alignment for different join types are: + // + // 1) Right & FullJoin -- to append all missing probe-side indices between + // previous (excluding) and current joined indices. + // 2) SemiJoin -- deduplicate probe indices in range between previous + // (excluding) and current joined indices. + // 3) AntiJoin -- return only missing indices in range between + // previous and current joined indices. + // Inclusion/exclusion of the indices themselves don't matter + // + // As a summary -- alignment range can be produced based only on + // joined (matched with filters applied) probe side indices, excluding starting one + // (left from previous iteration). + + // if any rows have been joined -- get last joined probe-side (right) row + // it's important that index counts as "joined" after hash collisions checks + // and join filters applied. + let last_joined_right_idx = match right_indices.len() { + 0 => None, + n => Some(right_indices.value(n - 1) as usize), + }; + + // Calculate range and perform alignment. + // In case probe batch has been processed -- align all remaining rows. + let index_alignment_range_start = state.joined_probe_idx.map_or(0, |v| v + 1); + let index_alignment_range_end = if next_offset.is_none() { + state.batch.num_rows() + } else { + last_joined_right_idx.map_or(0, |v| v + 1) + }; + + let (left_indices, right_indices) = adjust_indices_by_join_type( + left_indices, + right_indices, + index_alignment_range_start..index_alignment_range_end, + self.join_type, + self.right_side_ordered, + )?; + + // Build output batch and push to coalescer + let (build_batch, probe_batch, join_side) = + if self.join_type == JoinType::RightMark { + (&state.batch, build_side.left_data.batch(), JoinSide::Right) + } else { + (build_side.left_data.batch(), &state.batch, JoinSide::Left) + }; + + let batch = build_batch_from_indices( + &self.schema, + build_batch, + probe_batch, + &left_indices, + &right_indices, + &self.column_indices, + join_side, + self.join_type, + )?; + + let push_status = self.output_buffer.push_batch(batch)?; + + timer.done(); + + // If limit reached, finish and move to Completed state + if push_status == PushBatchStatus::LimitReached { + self.output_buffer.finish()?; + self.state = HashJoinStreamState::Completed; + return Ok(StatefulStreamResult::Continue); + } + + if next_offset.is_none() { + self.state = HashJoinStreamState::FetchProbeBatch; + } else { + state.advance( + next_offset + .ok_or_else(|| internal_datafusion_err!("unexpected None offset"))?, + last_joined_right_idx, + ) + }; + + Ok(StatefulStreamResult::Continue) + } + + /// Processes unmatched build-side rows for certain join types and produces output batch + /// + /// Updates state to `Completed` + fn process_unmatched_build_batch( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + let timer = self.join_metrics.join_time.timer(); + + if !need_produce_result_in_final(self.join_type) { + self.state = HashJoinStreamState::Completed; + return Ok(StatefulStreamResult::Continue); + } + + let build_side = self.build_side.try_as_ready()?; + + // For null-aware anti join, if probe side had NULL, no rows should be output + // Check shared atomic state to get global knowledge across all partitions + if self.null_aware + && build_side + .left_data + .probe_side_has_null + .load(Ordering::Relaxed) + { + timer.done(); + self.state = HashJoinStreamState::Completed; + return Ok(StatefulStreamResult::Continue); + } + if !build_side.left_data.report_probe_completed() { + self.state = HashJoinStreamState::Completed; + return Ok(StatefulStreamResult::Continue); + } + + // use the global left bitmap to produce the left indices and right indices + let (mut left_side, mut right_side) = get_final_indices_from_shared_bitmap( + build_side.left_data.visited_indices_bitmap(), + self.join_type, + true, + ); + + // For null-aware anti join, filter out LEFT rows with NULL in join keys + // BUT only if the probe side (RIGHT) was non-empty. If probe side is empty, + // NULL NOT IN (empty) = TRUE, so NULL rows should be returned. + // Use shared atomic state to get global knowledge across all partitions + if self.null_aware + && self.join_type == JoinType::LeftAnti + && build_side + .left_data + .probe_side_non_empty + .load(Ordering::Relaxed) + { + // Since null_aware validation ensures single column join, we only check the first column + let build_key_column = &build_side.left_data.values()[0]; + + // Filter out indices where the key is NULL + let filtered_indices: Vec<u64> = left_side + .iter() + .filter_map(|idx| { + let idx_usize = idx.unwrap() as usize; + if build_key_column.is_null(idx_usize) { + None // Skip rows with NULL keys + } else { + Some(idx.unwrap()) + } + }) + .collect(); + + left_side = UInt64Array::from(filtered_indices); + + // Update right_side to match the new length + let mut builder = arrow::array::UInt32Builder::with_capacity(left_side.len()); + builder.append_nulls(left_side.len()); + right_side = builder.finish(); + } + + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(left_side.len()); + + timer.done(); + + self.state = HashJoinStreamState::Completed; + + // Push final unmatched indices to output buffer + if !left_side.is_empty() { + let empty_right_batch = RecordBatch::new_empty(self.right.schema()); + let batch = build_batch_from_indices( + &self.schema, + build_side.left_data.batch(), + &empty_right_batch, + &left_side, + &right_side, + &self.column_indices, + JoinSide::Left, + self.join_type, + )?; + let push_status = self.output_buffer.push_batch(batch)?; + + // If limit reached, finish the coalescer + if push_status == PushBatchStatus::LimitReached { + self.output_buffer.finish()?; + } + } + + Ok(StatefulStreamResult::Continue) + } +} + +impl Stream for HashJoinStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_impl(cx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::joins::hash_join::shared_bounds::{ + PushdownStrategy, completed_partitions_for_test, + make_partitioned_accumulator_for_test, + }; + + fn empty_build_data(partition_id: usize) -> PartitionBuildData { + PartitionBuildData::Partitioned { + partition_id, + pushdown: PushdownStrategy::Empty, + bounds: PartitionBounds::new(vec![]), + keys_have_null: false, + } + } + + fn partitioned_handle(acc: &Arc<SharedBuildAccumulator>) -> BuildReportHandle { + BuildReportHandle::new(0, PartitionMode::Partitioned, Some(Arc::clone(acc))) + } + + #[test] + fn build_report_handle_cancels_scheduled_partition_on_drop() { + let acc = Arc::new(make_partitioned_accumulator_for_test(2)); + + { + let mut handle = partitioned_handle(&acc); + handle.schedule(empty_build_data(0)); + assert_eq!(handle.state(), &BuildReportState::Scheduled); + } + + assert_eq!(completed_partitions_for_test(&acc), 1); + } + + #[test] + fn build_report_handle_does_not_cancel_delivered_partition_on_drop() { + let acc = Arc::new(make_partitioned_accumulator_for_test(1)); + + { + let mut handle = partitioned_handle(&acc); + handle.schedule(empty_build_data(0)); + let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref()); + assert!(matches!(handle.poll_delivery(&mut cx), Poll::Ready(Ok(())))); + assert_eq!(handle.state(), &BuildReportState::Delivered); + } + + assert_eq!(completed_partitions_for_test(&acc), 1); + } + + #[test] + fn build_report_handle_cancel_pending_is_idempotent() { + let acc = Arc::new(make_partitioned_accumulator_for_test(2)); + let mut handle = partitioned_handle(&acc); + handle.schedule(empty_build_data(0)); + + handle.cancel_pending(); + handle.cancel_pending(); + + assert_eq!(handle.state(), &BuildReportState::Canceled); + assert_eq!(completed_partitions_for_test(&acc), 1); + } + + #[test] + fn build_report_handle_no_accumulator_finalizes() { + let mut handle = BuildReportHandle::new(0, PartitionMode::Partitioned, None); + + handle.schedule(empty_build_data(0)); + handle.cancel_pending(); + + assert_eq!(handle.state(), &BuildReportState::Finalized); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/join_filter.rs b/native/vendor/datafusion-physical-plan/src/joins/join_filter.rs new file mode 100644 index 0000000..de5df2b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/join_filter.rs
@@ -0,0 +1,108 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::joins::utils::ColumnIndex; +use arrow::datatypes::SchemaRef; +use datafusion_common::JoinSide; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use std::{fmt::Display, sync::Arc}; + +/// Filter applied before join output. Fields are crate-public to allow +/// downstream implementations to experiment with custom joins. +#[derive(Debug, Clone)] +pub struct JoinFilter { + /// Filter expression + pub(crate) expression: Arc<dyn PhysicalExpr>, + /// Column indices required to construct intermediate batch for filtering + pub(crate) column_indices: Vec<ColumnIndex>, + /// Physical schema of intermediate batch + pub(crate) schema: SchemaRef, +} + +/// For display in `EXPLAIN` plans, only expression with column names is needed, +/// it output expression like `(col1 + col2) = 0` +impl Display for JoinFilter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.expression.fmt_sql(f) + } +} + +impl JoinFilter { + /// Creates new JoinFilter + pub fn new( + expression: Arc<dyn PhysicalExpr>, + column_indices: Vec<ColumnIndex>, + schema: SchemaRef, + ) -> JoinFilter { + JoinFilter { + expression, + column_indices, + schema, + } + } + + /// Helper for building ColumnIndex vector from left and right indices + pub fn build_column_indices( + left_indices: Vec<usize>, + right_indices: Vec<usize>, + ) -> Vec<ColumnIndex> { + left_indices + .into_iter() + .map(|i| ColumnIndex { + index: i, + side: JoinSide::Left, + }) + .chain(right_indices.into_iter().map(|i| ColumnIndex { + index: i, + side: JoinSide::Right, + })) + .collect() + } + + /// Filter expression + pub fn expression(&self) -> &Arc<dyn PhysicalExpr> { + &self.expression + } + + /// Column indices for intermediate batch creation + pub fn column_indices(&self) -> &[ColumnIndex] { + &self.column_indices + } + + /// Intermediate batch schema + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Rewrites the join filter if the inputs to the join are rewritten + pub fn swap(&self) -> JoinFilter { + let column_indices = self + .column_indices() + .iter() + .map(|idx| ColumnIndex { + index: idx.index, + side: idx.side.negate(), + }) + .collect(); + + JoinFilter::new( + Arc::clone(self.expression()), + column_indices, + Arc::clone(self.schema()), + ) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/join_hash_map.rs b/native/vendor/datafusion-physical-plan/src/joins/join_hash_map.rs new file mode 100644 index 0000000..454cc91 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/join_hash_map.rs
@@ -0,0 +1,572 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This file contains the implementation of the `JoinHashMap` struct, which +//! is used to store the mapping between hash values based on the build side +//! ["on" values] to a list of indices with this key's value. + +use std::fmt::{self, Debug}; +use std::ops::Sub; + +use arrow::array::BooleanArray; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::datatypes::ArrowNativeType; +use hashbrown::HashTable; +use hashbrown::hash_table::Entry::{Occupied, Vacant}; + +/// Maps a `u64` hash value based on the build side ["on" values] to a list of indices with this key's value. +/// +/// By allocating a `HashMap` with capacity for *at least* the number of rows for entries at the build side, +/// we make sure that we don't have to re-hash the hashmap, which needs access to the key (the hash in this case) value. +/// +/// E.g. 1 -> [3, 6, 8] indicates that the column values map to rows 3, 6 and 8 for hash value 1 +/// As the key is a hash value, we need to check possible hash collisions in the probe stage +/// During this stage it might be the case that a row is contained the same hashmap value, +/// but the values don't match. Those are checked in the `equal_rows_arr` method. +/// +/// The indices (values) are stored in a separate chained list stored as `Vec<u32>` or `Vec<u64>`. +/// +/// The first value (+1) is stored in the hashmap, whereas the next value is stored in array at the position value. +/// +/// The chain can be followed until the value "0" has been reached, meaning the end of the list. +/// Also see chapter 5.3 of [Balancing vectorized query execution with bandwidth-optimized storage](https://dare.uva.nl/search?identifier=5ccbb60a-38b8-4eeb-858a-e7735dd37487) +/// +/// # Example +/// +/// ``` text +/// See the example below: +/// +/// Insert (10,1) <-- insert hash value 10 with row index 1 +/// map: +/// ---------- +/// | 10 | 2 | +/// ---------- +/// next: +/// --------------------- +/// | 0 | 0 | 0 | 0 | 0 | +/// --------------------- +/// Insert (20,2) +/// map: +/// ---------- +/// | 10 | 2 | +/// | 20 | 3 | +/// ---------- +/// next: +/// --------------------- +/// | 0 | 0 | 0 | 0 | 0 | +/// --------------------- +/// Insert (10,3) <-- collision! row index 3 has a hash value of 10 as well +/// map: +/// ---------- +/// | 10 | 4 | +/// | 20 | 3 | +/// ---------- +/// next: +/// --------------------- +/// | 0 | 0 | 0 | 2 | 0 | <--- hash value 10 maps to 4,2 (which means indices values 3,1) +/// --------------------- +/// Insert (10,4) <-- another collision! row index 4 ALSO has a hash value of 10 +/// map: +/// --------- +/// | 10 | 5 | +/// | 20 | 3 | +/// --------- +/// next: +/// --------------------- +/// | 0 | 0 | 0 | 2 | 4 | <--- hash value 10 maps to 5,4,2 (which means indices values 4,3,1) +/// --------------------- +/// ``` +/// +/// Here we have an option between creating a `JoinHashMapType` using `u32` or `u64` indices +/// based on how many rows were being used for indices. +/// +/// At runtime we choose between using `JoinHashMapU32` and `JoinHashMapU64` which oth implement +/// `JoinHashMapType`. +/// +/// ## Note on use of this trait as a public API +/// This is currently a public trait but is mainly intended for internal use within DataFusion. +/// For example, we may compare references to `JoinHashMapType` implementations by pointer equality +/// rather than deep equality of contents, as deep equality would be expensive and in our usage +/// patterns it is impossible for two different hash maps to have identical contents in a practical sense. +pub trait JoinHashMapType: Send + Sync { + fn extend_zero(&mut self, len: usize); + + fn update_from_iter<'a>( + &mut self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + Send + 'a>, + deleted_offset: usize, + ); + + fn get_matched_indices<'a>( + &self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + 'a>, + deleted_offset: Option<usize>, + ) -> (Vec<u32>, Vec<u64>); + + /// Probe rows marked NULL in `valid_keys` are skipped without a lookup: + /// their key contains a NULL, which cannot match any build row under + /// `NullEquality::NullEqualsNothing`. Pass `None` when every probe key is + /// matchable. + fn get_matched_indices_with_limit_offset( + &self, + hash_values: &[u64], + valid_keys: Option<&NullBuffer>, + limit: usize, + offset: MapOffset, + input_indices: &mut Vec<u32>, + match_indices: &mut Vec<u64>, + ) -> Option<MapOffset>; + + /// Returns a BooleanArray indicating which of the provided hashes exist in the map. + fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray; + + /// Returns `true` if the join hash map contains no entries. + fn is_empty(&self) -> bool; + + /// Returns the number of entries in the join hash map. + fn len(&self) -> usize; +} + +pub struct JoinHashMapU32 { + // Stores hash value to last row index + map: HashTable<(u64, u32)>, + // Stores indices in chained list data structure + next: Vec<u32>, +} + +impl JoinHashMapU32 { + #[cfg(test)] + pub(crate) fn new(map: HashTable<(u64, u32)>, next: Vec<u32>) -> Self { + Self { map, next } + } + + pub fn with_capacity(cap: usize) -> Self { + Self { + map: HashTable::with_capacity(cap), + next: vec![0; cap], + } + } +} + +impl Debug for JoinHashMapU32 { + fn fmt(&self, _f: &mut fmt::Formatter) -> fmt::Result { + Ok(()) + } +} + +impl JoinHashMapType for JoinHashMapU32 { + fn extend_zero(&mut self, _: usize) {} + + fn update_from_iter<'a>( + &mut self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + Send + 'a>, + deleted_offset: usize, + ) { + update_from_iter::<u32>(&mut self.map, &mut self.next, iter, deleted_offset); + } + + fn get_matched_indices<'a>( + &self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + 'a>, + deleted_offset: Option<usize>, + ) -> (Vec<u32>, Vec<u64>) { + get_matched_indices::<u32>(&self.map, &self.next, iter, deleted_offset) + } + + fn get_matched_indices_with_limit_offset( + &self, + hash_values: &[u64], + valid_keys: Option<&NullBuffer>, + limit: usize, + offset: MapOffset, + input_indices: &mut Vec<u32>, + match_indices: &mut Vec<u64>, + ) -> Option<MapOffset> { + get_matched_indices_with_limit_offset::<u32>( + &self.map, + &self.next, + hash_values, + valid_keys, + limit, + offset, + input_indices, + match_indices, + ) + } + + fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray { + contain_hashes(&self.map, hash_values) + } + + fn is_empty(&self) -> bool { + self.map.is_empty() + } + + fn len(&self) -> usize { + self.map.len() + } +} + +pub struct JoinHashMapU64 { + // Stores hash value to last row index + map: HashTable<(u64, u64)>, + // Stores indices in chained list data structure + next: Vec<u64>, +} + +impl JoinHashMapU64 { + #[cfg(test)] + pub(crate) fn new(map: HashTable<(u64, u64)>, next: Vec<u64>) -> Self { + Self { map, next } + } + + pub fn with_capacity(cap: usize) -> Self { + Self { + map: HashTable::with_capacity(cap), + next: vec![0; cap], + } + } +} + +impl Debug for JoinHashMapU64 { + fn fmt(&self, _f: &mut fmt::Formatter) -> fmt::Result { + Ok(()) + } +} + +impl JoinHashMapType for JoinHashMapU64 { + fn extend_zero(&mut self, _: usize) {} + + fn update_from_iter<'a>( + &mut self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + Send + 'a>, + deleted_offset: usize, + ) { + update_from_iter::<u64>(&mut self.map, &mut self.next, iter, deleted_offset); + } + + fn get_matched_indices<'a>( + &self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + 'a>, + deleted_offset: Option<usize>, + ) -> (Vec<u32>, Vec<u64>) { + get_matched_indices::<u64>(&self.map, &self.next, iter, deleted_offset) + } + + fn get_matched_indices_with_limit_offset( + &self, + hash_values: &[u64], + valid_keys: Option<&NullBuffer>, + limit: usize, + offset: MapOffset, + input_indices: &mut Vec<u32>, + match_indices: &mut Vec<u64>, + ) -> Option<MapOffset> { + get_matched_indices_with_limit_offset::<u64>( + &self.map, + &self.next, + hash_values, + valid_keys, + limit, + offset, + input_indices, + match_indices, + ) + } + + fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray { + contain_hashes(&self.map, hash_values) + } + + fn is_empty(&self) -> bool { + self.map.is_empty() + } + + fn len(&self) -> usize { + self.map.len() + } +} + +use crate::joins::MapOffset; +use crate::joins::chain::traverse_chain; + +pub fn update_from_iter<'a, T>( + map: &mut HashTable<(u64, T)>, + next: &mut [T], + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + Send + 'a>, + deleted_offset: usize, +) where + T: Copy + TryFrom<usize> + PartialOrd, + <T as TryFrom<usize>>::Error: Debug, +{ + for (row, &hash_value) in iter { + let entry = map.entry( + hash_value, + |&(hash, _)| hash_value == hash, + |&(hash, _)| hash, + ); + + match entry { + Occupied(mut occupied_entry) => { + // Already exists: add index to next array + let (_, index) = occupied_entry.get_mut(); + let prev_index = *index; + // Store new value inside hashmap + *index = T::try_from(row + 1).unwrap(); + // Update chained Vec at `row` with previous value + next[row - deleted_offset] = prev_index; + } + Vacant(vacant_entry) => { + vacant_entry.insert((hash_value, T::try_from(row + 1).unwrap())); + } + } + } +} + +pub fn get_matched_indices<'a, T>( + map: &HashTable<(u64, T)>, + next: &[T], + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + 'a>, + deleted_offset: Option<usize>, +) -> (Vec<u32>, Vec<u64>) +where + T: Copy + TryFrom<usize> + PartialOrd + Into<u64> + Sub<Output = T>, + <T as TryFrom<usize>>::Error: Debug, +{ + let mut input_indices = vec![]; + let mut match_indices = vec![]; + let zero = T::try_from(0).unwrap(); + let one = T::try_from(1).unwrap(); + + for (row_idx, hash_value) in iter { + // Get the hash and find it in the index + if let Some((_, index)) = map.find(*hash_value, |(hash, _)| *hash_value == *hash) + { + let mut i = *index - one; + loop { + let match_row_idx = if let Some(offset) = deleted_offset { + let offset = T::try_from(offset).unwrap(); + // This arguments means that we prune the next index way before here. + if i < offset { + // End of the list due to pruning + break; + } + i - offset + } else { + i + }; + match_indices.push(match_row_idx.into()); + input_indices.push(row_idx as u32); + // Follow the chain to get the next index value + let next_chain = next[match_row_idx.into() as usize]; + if next_chain == zero { + // end of list + break; + } + i = next_chain - one; + } + } + } + + (input_indices, match_indices) +} + +#[expect(clippy::too_many_arguments)] +pub fn get_matched_indices_with_limit_offset<T>( + map: &HashTable<(u64, T)>, + next_chain: &[T], + hash_values: &[u64], + valid_keys: Option<&NullBuffer>, + limit: usize, + offset: MapOffset, + input_indices: &mut Vec<u32>, + match_indices: &mut Vec<u64>, +) -> Option<MapOffset> +where + T: Copy + TryFrom<usize> + PartialOrd + Into<u64> + Sub<Output = T>, + <T as TryFrom<usize>>::Error: Debug, + T: ArrowNativeType, +{ + // Clear the buffer before producing new results + input_indices.clear(); + match_indices.clear(); + let one = T::try_from(1).unwrap(); + + // Check if hashmap consists of unique values + // If so, we can skip the chain traversal + if map.len() == next_chain.len() { + let start = offset.0; + let end = (start + limit).min(hash_values.len()); + for (i, &hash) in hash_values[start..end].iter().enumerate() { + // NULL keys cannot match any build row + if valid_keys.is_some_and(|valid| valid.is_null(start + i)) { + continue; + } + if let Some((_, idx)) = map.find(hash, |(h, _)| hash == *h) { + input_indices.push(start as u32 + i as u32); + match_indices.push((*idx - one).into()); + } + } + return if end == hash_values.len() { + None + } else { + Some((end, None)) + }; + } + + let mut remaining_output = limit; + + // Calculate initial `hash_values` index before iterating + let to_skip = match offset { + // None `initial_next_idx` indicates that `initial_idx` processing hasn't been started + (idx, None) => idx, + // Zero `initial_next_idx` indicates that `initial_idx` has been processed during + // previous iteration, and it should be skipped + (idx, Some(0)) => idx + 1, + // Otherwise, process remaining `initial_idx` matches by traversing `next_chain`, + // to start with the next index + (idx, Some(next_idx)) => { + let next_idx: T = T::usize_as(next_idx as usize); + let is_last = idx == hash_values.len() - 1; + if let Some(next_offset) = traverse_chain( + next_chain, + idx, + next_idx, + &mut remaining_output, + input_indices, + match_indices, + is_last, + ) { + return Some(next_offset); + } + idx + 1 + } + }; + + let hash_values_len = hash_values.len(); + for (i, &hash) in hash_values[to_skip..].iter().enumerate() { + let row_idx = to_skip + i; + // NULL keys cannot match any build row + if valid_keys.is_some_and(|valid| valid.is_null(row_idx)) { + continue; + } + if let Some((_, idx)) = map.find(hash, |(h, _)| hash == *h) { + let idx: T = *idx; + let is_last = row_idx == hash_values_len - 1; + if let Some(next_offset) = traverse_chain( + next_chain, + row_idx, + idx, + &mut remaining_output, + input_indices, + match_indices, + is_last, + ) { + return Some(next_offset); + } + } + } + None +} + +pub fn contain_hashes<T>(map: &HashTable<(u64, T)>, hash_values: &[u64]) -> BooleanArray { + let buffer = BooleanBuffer::collect_bool(hash_values.len(), |i| { + let hash = hash_values[i]; + map.find(hash, |(h, _)| hash == *h).is_some() + }); + BooleanArray::new(buffer, None) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_contain_hashes() { + let mut hash_map = JoinHashMapU32::with_capacity(10); + hash_map.update_from_iter(Box::new([10u64, 20u64, 30u64].iter().enumerate()), 0); + + let probe_hashes = vec![10, 11, 20, 21, 30, 31]; + let array = hash_map.contain_hashes(&probe_hashes); + + assert_eq!(array.len(), probe_hashes.len()); + + for (i, &hash) in probe_hashes.iter().enumerate() { + if matches!(hash, 10 | 20 | 30) { + assert!(array.value(i), "Hash {hash} should exist in the map"); + } else { + assert!(!array.value(i), "Hash {hash} should NOT exist in the map"); + } + } + } + + #[test] + fn test_get_matched_indices_skips_invalid_keys() { + let mut hash_map = JoinHashMapU32::with_capacity(3); + hash_map.update_from_iter(Box::new([10u64, 20u64, 30u64].iter().enumerate()), 0); + + let probe_hashes = vec![10, 20, 30]; + // The probe row for hash 20 has a NULL key and must not match. + let valid_keys = NullBuffer::from(vec![true, false, true]); + + let mut input_indices = vec![]; + let mut match_indices = vec![]; + let next_offset = hash_map.get_matched_indices_with_limit_offset( + &probe_hashes, + Some(&valid_keys), + 8192, + (0, None), + &mut input_indices, + &mut match_indices, + ); + + assert_eq!(next_offset, None); + assert_eq!(input_indices, vec![0, 2]); + assert_eq!(match_indices, vec![0, 2]); + } + + #[test] + fn test_get_matched_indices_skips_invalid_keys_with_duplicates() { + // Duplicate build keys chain multiple rows under one hash value. + let mut hash_map = JoinHashMapU32::with_capacity(4); + hash_map.update_from_iter( + Box::new([10u64, 20u64, 10u64, 20u64].iter().enumerate()), + 0, + ); + + let probe_hashes = vec![10, 20]; + // The probe row for hash 10 has a NULL key: none of the build rows in + // its chain may match, while the valid probe row for hash 20 must + // still match its entire chain. + let valid_keys = NullBuffer::from(vec![false, true]); + + let mut input_indices = vec![]; + let mut match_indices = vec![]; + let next_offset = hash_map.get_matched_indices_with_limit_offset( + &probe_hashes, + Some(&valid_keys), + 8192, + (0, None), + &mut input_indices, + &mut match_indices, + ); + + assert_eq!(next_offset, None); + assert_eq!(input_indices, vec![1, 1]); + assert_eq!(match_indices, vec![3, 1]); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/key_range_bitmap.rs b/native/vendor/datafusion-physical-plan/src/joins/key_range_bitmap.rs new file mode 100644 index 0000000..527ee76 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/key_range_bitmap.rs
@@ -0,0 +1,238 @@ +// 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. + +//! Bitmap over a hash join build side's key range, for container pruning. + +use arrow::array::{Array, ArrayRef, BooleanArray, downcast_integer_array}; +use arrow::buffer::{BooleanBuffer, MutableBuffer}; +use arrow::util::bit_util; +use datafusion_common::ScalarValue; +use num_traits::AsPrimitive; + +use crate::joins::array_map::ArrayMap; + +/// Largest bitmap, in bits (128 KiB), which bounds what one build side can hold. +const MAX_BUCKETS: u64 = 1 << 20; + +/// Which parts of a build side's key range hold at least one key. +/// +/// Keys map order-preservingly onto buckets spanning `[min, max]`, one bit each, +/// so "can a container holding `[lo, hi]` match?" is a scan of the bits spanning +/// `[lo, hi]`: all clear proves no build key lies in it. +/// +/// Below, keys 3 to 97 over eight buckets: `offset` is 3, `span` is 97 - 3, and +/// eight buckets need `shift` = 4, so each covers `1 << 4` keys starting at +/// `offset`. Only buckets 0 to `span >> shift` can hold a key. Real bitmaps run +/// to `MAX_BUCKETS`; eight is what fits on a line. +/// +/// ```text +/// keys 3 12 42 97 +/// v v v v +/// bucket | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | +/// covers 3-18 19-34 35-50 51-66 67-82 83-98 unused unused +/// bit 1 0 1 0 0 1 0 0 +/// ``` +/// +/// A container is kept when any bucket its `[min, max]` touches has a bit set. +#[derive(Debug)] +pub struct KeyRangeBitmap { + /// Smallest key, in the `u64` ordering the join's array map uses. + offset: u64, + /// `max - min`: a key is in range iff `key.wrapping_sub(offset) <= span`. + span: u64, + /// Each bucket covers `1 << shift` keys. + shift: u32, + /// One bit per bucket. + bits: BooleanBuffer, +} + +impl KeyRangeBitmap { + /// Maps `array`'s non-null values, which must lie within `[min, max]`. + /// + /// `None` when the key type has no `u64` ordering, or when the bitmap could + /// not exclude anything the `[min, max]` bounds predicate does not already. + pub fn try_new( + array: &ArrayRef, + min: &ScalarValue, + max: &ScalarValue, + distinct_keys: usize, + ) -> Option<Self> { + if !ArrayMap::is_supported_type(array.data_type()) { + return None; + } + let offset = ArrayMap::key_to_u64(min)?; + let span = ArrayMap::key_to_u64(max)?.wrapping_sub(offset); + + // One bucket per key value where the range allows it, so the bitmap is + // exact; otherwise as many as the cap permits. + let buckets = span.saturating_add(1).min(MAX_BUCKETS).next_power_of_two(); + + // Use a heuristic to bail-out for the common shape which renders each bucket + // set (contiguous surrogate keys, a date range) before actually computing it. + if (distinct_keys as u64).saturating_mul(2) > span.min(buckets) { + return None; + } + + let shift = (0u32..63).find(|s| (span >> s) < buckets).unwrap_or(63); + + let bits = bucket_bits(array, offset, shift, buckets); + + // Every bucket set means nothing can be excluded. + let reachable = ((span >> shift) + 1) as usize; + bits.slice(0, reachable).has_false().then_some(Self { + offset, + span, + shift, + bits, + }) + } + + /// One Boolean per `[min[i], max[i]]` container interval: `false` proves no + /// build key lies in that container, `None` where a bound is absent or the + /// statistics are not integer keys. + pub fn may_contain_ranges(&self, min: &dyn Array, max: &dyn Array) -> BooleanArray { + // Both bounds must be the same integer type for the tuple pattern to match. + downcast_integer_array!( + (min, max) => { + min.iter() + .zip(max.iter()) + .map(|bounds| match bounds { + (Some(lo), Some(hi)) => { + Some(self.may_contain_range(lo.as_(), hi.as_())) + } + _ => None, + }) + .collect() + } + _ => BooleanArray::new_null(min.len()), + ) + } + + /// Might a container spanning `[lo, hi]` hold a build key? `false` proves it + /// cannot. Bounds use the same `u64` key ordering as the join's array map. + pub fn may_contain_range(&self, lo: u64, hi: u64) -> bool { + // Offsets from the minimum; anything outside `[min, max]` wraps to + // something huge, which the clamping below turns into "past the end". + let lo_off = lo.wrapping_sub(self.offset); + let hi_off = hi.wrapping_sub(self.offset); + if lo_off > self.span && hi_off > self.span && lo_off <= hi_off { + return false; + } + let lo_b = if lo_off > self.span { + 0 + } else { + lo_off >> self.shift + }; + let hi_b = (hi_off.min(self.span) >> self.shift).max(lo_b); + self.bits + .slice(lo_b as usize, (hi_b - lo_b + 1) as usize) + .has_true() + } + + /// Bytes held, for the build-side memory reservation. + pub fn size(&self) -> usize { + self.bits.inner().capacity() + } + + /// Buckets set and buckets total, for `EXPLAIN`. + pub fn fill_stats(&self) -> (usize, usize) { + (self.bits.count_set_bits(), self.bits.len()) + } +} + +/// One bit per bucket, set for each bucket holding a non-null key of `keys`. +fn bucket_bits(keys: &dyn Array, offset: u64, shift: u32, buckets: u64) -> BooleanBuffer { + let mut words = MutableBuffer::new_null(buckets as usize); + downcast_integer_array!( + keys => { + let bytes = words.as_slice_mut(); + let mut set = |key: u64| { + let bucket = (key.wrapping_sub(offset) >> shift).min(buckets - 1); + bit_util::set_bit(bytes, bucket as usize); + }; + if keys.null_count() == 0 { + keys.values().iter().for_each(|key| set((*key).as_())); + } else { + keys.iter().flatten().for_each(|key| set(key.as_())); + } + } + _ => unreachable!("guarded by ArrayMap::is_supported_type"), + ); + BooleanBuffer::new(words.into(), 0, buckets as usize) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use std::sync::Arc; + + fn bitmap(keys: &[i64]) -> Option<KeyRangeBitmap> { + let a: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); + KeyRangeBitmap::try_new( + &a, + &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), + &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), + keys.iter().collect::<std::collections::HashSet<_>>().len(), + ) + } + fn may(b: &KeyRangeBitmap, lo: i64, hi: i64) -> bool { + b.may_contain_range(lo as u64, hi as u64) + } + + #[test] + fn prunes_gaps_and_never_drops_a_match() { + let keys: Vec<i64> = (0..2_000_000).filter(|k| k % 10000 < 200).collect(); + let s = bitmap(&keys).expect("not saturated"); + let mut kept = 0; + for i in 0..2000i64 { + let (lo, hi) = (1000 * i, 1000 * i + 999); + let truth = keys + .binary_search(&lo) + .map_or_else(|p| keys.get(p).is_some_and(|k| *k <= hi), |_| true); + let got = may(&s, lo, hi); + assert!( + !truth || got, + "dropped a container holding a key: {lo}..{hi}" + ); + kept += got as usize; + } + assert!(kept <= 220, "expected ~200 kept, got {kept}"); + } + + #[test] + fn contiguous_keys_no_bitmap() { + assert!(bitmap(&(0..5000i64).collect::<Vec<_>>()).is_none()); + } + + #[test] + fn test_bitmap_range_probing() { + let s = bitmap(&[0, 1_000_000]).expect("not saturated"); + assert!(!may(&s, -500, -10)); + assert!(!may(&s, 2_000_000, 3_000_000)); + assert!(may(&s, -500, 10)); + + let s = bitmap(&[-1_000_000, -5, 0, 5, 1_000_000]).expect("not saturated"); + assert!(may(&s, -10, 10)); + assert!(!may(&s, -900_000, -800_000)); + + // The widest possible span must not overflow the bucket sizing. + let s = bitmap(&[i64::MIN, 0, i64::MAX]).expect("not saturated"); + assert!(may(&s, -1, 1)); + assert!(!may(&s, 1 << 50, 1 << 51)); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/mod.rs b/native/vendor/datafusion-physical-plan/src/joins/mod.rs new file mode 100644 index 0000000..ed9e1ad --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/mod.rs
@@ -0,0 +1,118 @@ +// 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. + +//! DataFusion Join implementations + +use arrow::array::BooleanBufferBuilder; +pub use cross_join::CrossJoinExec; +use datafusion_physical_expr::PhysicalExprRef; +pub use hash_join::{ + HashExpr, HashJoinExec, HashJoinExecBuilder, HashTableLookupExpr, SeededRandomState, +}; +pub use nested_loop_join::{NestedLoopJoinExec, NestedLoopJoinExecBuilder}; +use parking_lot::Mutex; +// Note: SortMergeJoin is not used in plans yet +pub use piecewise_merge_join::PiecewiseMergeJoinExec; +pub use sort_merge_join::SortMergeJoinExec; +pub use symmetric_hash_join::SymmetricHashJoinExec; +pub mod chain; +mod cross_join; +mod hash_join; +mod nested_loop_join; +mod piecewise_merge_join; +#[cfg(feature = "proto")] +mod proto; +mod sort_merge_join; +mod stream_join_utils; +mod symmetric_hash_join; +pub mod utils; + +mod array_map; +mod join_filter; +/// Hash map implementations for join operations. +/// +/// Note: This module is public for internal testing purposes only +/// and is not guaranteed to be stable across versions. +pub mod join_hash_map; +pub mod key_range_bitmap; + +use array_map::ArrayMap; +use utils::JoinHashMapType; + +/// The build-side map of a hash join, indexing build rows by join key. +/// +/// Under [`NullEquality::NullEqualsNothing`], build rows with a NULL in any +/// join key column can never match a probe row and are omitted from the map. +/// [`Map::is_empty`] and [`Map::num_of_distinct_key`] therefore reflect the +/// *matchable* build rows: the map can be empty even when the build side +/// contains rows. +/// +/// [`NullEquality::NullEqualsNothing`]: datafusion_common::NullEquality::NullEqualsNothing +pub enum Map { + HashMap(Box<dyn JoinHashMapType>), + ArrayMap(ArrayMap), +} + +impl Map { + /// Returns the number of elements in the map. + pub fn num_of_distinct_key(&self) -> usize { + match self { + Map::HashMap(map) => map.len(), + Map::ArrayMap(array_map) => array_map.num_of_distinct_key(), + } + } + + /// Returns `true` if the map contains no elements. + pub fn is_empty(&self) -> bool { + self.num_of_distinct_key() == 0 + } +} + +pub(crate) type MapOffset = (usize, Option<u64>); + +#[cfg(test)] +pub mod test_utils; + +/// The on clause of the join, as vector of (left, right) columns. +pub type JoinOn = Vec<(PhysicalExprRef, PhysicalExprRef)>; +/// Reference for JoinOn. +pub type JoinOnRef<'a> = &'a [(PhysicalExprRef, PhysicalExprRef)]; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +/// Hash join Partitioning mode +pub enum PartitionMode { + /// Left/right children are partitioned using the left and right keys + Partitioned, + /// Left side will collected into one partition + CollectLeft, + /// DataFusion optimizer decides which PartitionMode + /// mode(Partitioned/CollectLeft) is optimal based on statistics. It will + /// also consider swapping the left and right inputs for the Join + Auto, +} + +/// Partitioning mode to use for symmetric hash join +#[derive(Hash, Clone, Copy, Debug, PartialEq, Eq)] +pub enum StreamJoinPartitionMode { + /// Left/right children are partitioned using the left and right keys + Partitioned, + /// Both sides will collected into one partition + SinglePartition, +} + +/// Shared bitmap for visited left-side indices +type SharedBitmapBuilder = Mutex<BooleanBufferBuilder>;
diff --git a/native/vendor/datafusion-physical-plan/src/joins/nested_loop_join.rs b/native/vendor/datafusion-physical-plan/src/joins/nested_loop_join.rs new file mode 100644 index 0000000..eb1df63 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/nested_loop_join.rs
@@ -0,0 +1,4144 @@ +// 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. + +//! [`NestedLoopJoinExec`]: joins without equijoin (equality predicates). + +use std::fmt::Formatter; +use std::ops::{BitOr, ControlFlow}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::task::Poll; + +use super::utils::{ + asymmetric_join_output_partitioning, need_produce_result_in_final, + reorder_output_after_swap, swap_join_projection, +}; +use crate::common::can_project; +use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::joins::SharedBitmapBuilder; +use crate::joins::utils::{ + BuildProbeJoinMetrics, ColumnIndex, JoinFilter, OnceAsync, OnceFut, + build_join_schema, check_join_is_valid, estimate_join_statistics, + need_produce_right_in_final, +}; +use crate::metrics::{ + Count, ExecutionPlanMetricsSet, MetricBuilder, MetricType, MetricsSet, RatioMetrics, +}; +use crate::projection::{ + EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection, + try_pushdown_through_join_with_column_indices, +}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, validate_child_count, +}; + +use arrow::array::{ + Array, BooleanArray, BooleanBufferBuilder, RecordBatchOptions, UInt32Array, + UInt64Array, new_null_array, +}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::{ + BatchCoalescer, concat_batches, filter, filter_record_batch, not, take, +}; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use arrow_schema::DataType; +use datafusion_common::cast::as_boolean_array; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + JoinSide, NullEquality, Result, ScalarValue, Statistics, arrow_err, + assert_eq_or_internal_err, internal_datafusion_err, internal_err, project_schema, + unwrap_or_internal_err, +}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{SpillFile, TaskContext}; +use datafusion_expr::JoinType; +use datafusion_physical_expr::equivalence::{ + ProjectionMapping, join_equivalence_properties, +}; + +use datafusion_physical_expr::projection::{ProjectionRef, combine_projections}; +use futures::{Stream, StreamExt, TryStreamExt}; +use log::debug; +use parking_lot::Mutex; + +use crate::metrics::SpillMetrics; +use crate::spill::replayable_spill_input::ReplayableStreamSource; +use crate::spill::spill_manager::SpillManager; + +#[expect(rustdoc::private_intra_doc_links)] +/// NestedLoopJoinExec is a build-probe join operator designed for joins that +/// do not have equijoin keys in their `ON` clause. +/// +/// # Execution Flow +/// +/// ```text +/// Incoming right batch +/// Left Side Buffered Batches +/// ┌───────────┐ ┌───────────────┐ +/// │ ┌───────┐ │ │ │ +/// │ │ │ │ │ │ +/// Current Left Row ───▶│ ├───────├─┤──────────┐ │ │ +/// │ │ │ │ │ └───────────────┘ +/// │ │ │ │ │ │ +/// │ │ │ │ │ │ +/// │ └───────┘ │ │ │ +/// │ ┌───────┐ │ │ │ +/// │ │ │ │ │ ┌─────┘ +/// │ │ │ │ │ │ +/// │ │ │ │ │ │ +/// │ │ │ │ │ │ +/// │ │ │ │ │ │ +/// │ └───────┘ │ ▼ ▼ +/// │ ...... │ ┌──────────────────────┐ +/// │ │ │X (Cartesian Product) │ +/// │ │ └──────────┬───────────┘ +/// └───────────┘ │ +/// │ +/// ▼ +/// ┌───────┬───────────────┐ +/// │ │ │ +/// │ │ │ +/// │ │ │ +/// └───────┴───────────────┘ +/// Intermediate Batch +/// (For join predicate evaluation) +/// ``` +/// +/// The execution follows a two-phase design: +/// +/// ## 1. Buffering Left Input +/// - The operator eagerly buffers all left-side input batches into memory, +/// util a memory limit is reached. +/// Currently, an out-of-memory error will be thrown if all the left-side input batches +/// cannot fit into memory at once. +/// In the future, it's possible to make this case finish execution. (see +/// 'Memory-limited Execution' section) +/// - The rationale for buffering the left side is that scanning the right side +/// can be expensive (e.g., decoding Parquet files), so buffering more left +/// rows reduces the number of right-side scan passes required. +/// +/// ## 2. Probing Right Input +/// - Right-side input is streamed batch by batch. +/// - For each right-side batch: +/// - It evaluates the join filter against the full buffered left input. +/// This results in a Cartesian product between the right batch and each +/// left row -- with the join predicate/filter applied -- for each inner +/// loop iteration. +/// - Matched results are accumulated into an output buffer. (see more in +/// `Output Buffering Strategy` section) +/// - This process continues until all right-side input is consumed. +/// +/// # Producing unmatched build-side data +/// - For special join types like left/full joins, it's required to also output +/// unmatched pairs. During execution, bitmaps are kept for both left and right +/// sides of the input; they'll be handled by dedicated states in `NLJStream`. +/// - The final output of the left side unmatched rows is handled by a single +/// partition for simplicity, since it only counts a small portion of the +/// execution time. (e.g. if probe side has 10k rows, the final output of +/// unmatched build side only roughly counts for 1/10k of the total time) +/// +/// # Output Buffering Strategy +/// The operator uses an intermediate output buffer to accumulate results. Once +/// the output threshold is reached (currently set to the same value as +/// `batch_size` in the configuration), the results will be eagerly output. +/// +/// # Extra Notes +/// - The operator always considers the **left** side as the build (buffered) side. +/// Therefore, the physical optimizer should assign the smaller input to the left. +/// - The design try to minimize the intermediate data size to approximately +/// 1 batch, for better cache locality and memory efficiency. +/// +/// # Memory-limited Execution +/// When the memory budget is exceeded during left-side buffering, the operator +/// falls back to a multi-pass strategy: +/// 1. Buffer as many left rows as fit in memory (one "chunk") +/// 2. On the first pass, the right side is both processed and spilled to disk +/// 3. For each subsequent left chunk, the right side is re-read from the spill file +/// +/// The fallback is triggered automatically when the initial in-memory load +/// fails with `ResourcesExhausted` and disk spilling is available. Each +/// output partition independently re-executes the left child and manages +/// its own spill state. +/// +/// All join types are supported. For RIGHT/FULL/RIGHT SEMI/RIGHT ANTI/ +/// RIGHT MARK joins, a global right-side bitmap (indexed by right batch +/// sequence number) accumulates matches across all left chunks. After the +/// last left chunk is processed, the right side is replayed one more time +/// to emit unmatched right rows using the accumulated bitmap. +/// +/// Tracking issue: <https://github.com/apache/datafusion/issues/15760> +/// +/// # Clone / Shared State +/// Note this structure includes a [`OnceAsync`] that is used to coordinate the +/// loading of the left side with the processing in each output stream. +/// Therefore it can not be [`Clone`] +#[derive(Debug)] +pub struct NestedLoopJoinExec { + /// left side + pub(crate) left: Arc<dyn ExecutionPlan>, + /// right side + pub(crate) right: Arc<dyn ExecutionPlan>, + /// Filters which are applied while finding matching rows + pub(crate) filter: Option<JoinFilter>, + /// How the join is performed + pub(crate) join_type: JoinType, + /// The full concatenated schema of left and right children should be distinct from + /// the output schema of the operator + join_schema: SchemaRef, + /// Future that consumes left input and buffers it in memory + /// + /// This structure is *shared* across all output streams. + /// + /// Each output stream waits on the `OnceAsync` to signal the completion of + /// the build(left) side data, and buffer them all for later joining. + build_side_data: OnceAsync<JoinLeftData>, + /// Shared left-side spill data for OOM fallback. + /// + /// When `build_side_data` fails with OOM, the first partition to + /// initiate fallback spills the entire left side to disk. Other + /// partitions share the same spill file via this `OnceAsync`, + /// avoiding redundant re-execution of the left child. + left_spill_data: Arc<OnceAsync<LeftSpillData>>, + /// Information of index and left / right placement of columns + column_indices: Vec<ColumnIndex>, + /// Projection to apply to the output of the join + projection: Option<ProjectionRef>, + + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +/// Helps to build [`NestedLoopJoinExec`]. +pub struct NestedLoopJoinExecBuilder { + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: JoinType, + filter: Option<JoinFilter>, + projection: Option<ProjectionRef>, +} + +impl NestedLoopJoinExecBuilder { + /// Make a new [`NestedLoopJoinExecBuilder`]. + pub fn new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: JoinType, + ) -> Self { + Self { + left, + right, + join_type, + filter: None, + projection: None, + } + } + + /// Set projection from the vector. + pub fn with_projection(self, projection: Option<Vec<usize>>) -> Self { + self.with_projection_ref(projection.map(Into::into)) + } + + /// Set projection from the shared reference. + pub fn with_projection_ref(mut self, projection: Option<ProjectionRef>) -> Self { + self.projection = projection; + self + } + + /// Set optional filter. + pub fn with_filter(mut self, filter: Option<JoinFilter>) -> Self { + self.filter = filter; + self + } + + /// Build resulting execution plan. + pub fn build(self) -> Result<NestedLoopJoinExec> { + let Self { + left, + right, + join_type, + filter, + projection, + } = self; + + let left_schema = left.schema(); + let right_schema = right.schema(); + check_join_is_valid(&left_schema, &right_schema, &[])?; + let (join_schema, column_indices) = + build_join_schema(&left_schema, &right_schema, &join_type); + let join_schema = Arc::new(join_schema); + let cache = NestedLoopJoinExec::compute_properties( + &left, + &right, + &join_schema, + join_type, + projection.as_deref(), + )?; + Ok(NestedLoopJoinExec { + left, + right, + filter, + join_type, + join_schema, + build_side_data: Default::default(), + left_spill_data: Arc::new(OnceAsync::default()), + column_indices, + projection, + metrics: Default::default(), + cache: Arc::new(cache), + }) + } +} + +impl From<&NestedLoopJoinExec> for NestedLoopJoinExecBuilder { + fn from(exec: &NestedLoopJoinExec) -> Self { + Self { + left: Arc::clone(exec.left()), + right: Arc::clone(exec.right()), + join_type: exec.join_type, + filter: exec.filter.clone(), + projection: exec.projection.clone(), + } + } +} + +impl NestedLoopJoinExec { + /// Try to create a new [`NestedLoopJoinExec`] + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + filter: Option<JoinFilter>, + join_type: &JoinType, + projection: Option<Vec<usize>>, + ) -> Result<Self> { + NestedLoopJoinExecBuilder::new(left, right, *join_type) + .with_projection(projection) + .with_filter(filter) + .build() + } + + /// left side + pub fn left(&self) -> &Arc<dyn ExecutionPlan> { + &self.left + } + + /// right side + pub fn right(&self) -> &Arc<dyn ExecutionPlan> { + &self.right + } + + /// Filters applied before join output + pub fn filter(&self) -> Option<&JoinFilter> { + self.filter.as_ref() + } + + /// How the join is performed + pub fn join_type(&self) -> &JoinType { + &self.join_type + } + + pub fn projection(&self) -> &Option<ProjectionRef> { + &self.projection + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + schema: &SchemaRef, + join_type: JoinType, + projection: Option<&[usize]>, + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let mut eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &join_type, + Arc::clone(schema), + &Self::maintains_input_order(join_type), + None, + // No on columns in nested loop join + &[], + )?; + + let mut output_partitioning = + asymmetric_join_output_partitioning(left, right, &join_type)?; + + let emission_type = if left.boundedness().is_unbounded() { + EmissionType::Final + } else if right.pipeline_behavior() == EmissionType::Incremental { + match join_type { + // If we only need to generate matched rows from the probe side, + // we can emit rows incrementally. + JoinType::Inner + | JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::Right + | JoinType::RightAnti + | JoinType::RightMark => EmissionType::Incremental, + // If we need to generate unmatched rows from the *build side*, + // we need to emit them at the end. + JoinType::Left + | JoinType::LeftAnti + | JoinType::LeftMark + | JoinType::Full => EmissionType::Both, + } + } else { + right.pipeline_behavior() + }; + + if let Some(projection) = projection { + // construct a map from the input expressions to the output expression of the Projection + let projection_mapping = ProjectionMapping::from_indices(projection, schema)?; + let out_schema = project_schema(schema, Some(&projection))?; + output_partitioning = + output_partitioning.project(&projection_mapping, &eq_properties); + eq_properties = eq_properties.project(&projection_mapping, out_schema); + } + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + emission_type, + boundedness_from_children([left, right]), + )) + } + + /// This join implementation does not preserve the input order of either side. + fn maintains_input_order(_join_type: JoinType) -> Vec<bool> { + vec![false, false] + } + + pub fn contains_projection(&self) -> bool { + self.projection.is_some() + } + + pub fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> { + let projection = projection.map(Into::into); + // check if the projection is valid + can_project(&self.schema(), projection.as_deref())?; + let projection = + combine_projections(projection.as_ref(), self.projection.as_ref())?; + NestedLoopJoinExecBuilder::from(self) + .with_projection_ref(projection) + .build() + } + + /// Returns a new `ExecutionPlan` that runs NestedLoopsJoins with the left + /// and right inputs swapped. + /// + /// # Notes: + /// + /// This function should be called BEFORE inserting any repartitioning + /// operators on the join's children. Check [`super::HashJoinExec::swap_inputs`] + /// for more details. + pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> { + let left = self.left(); + let right = self.right(); + let new_join = NestedLoopJoinExec::try_new( + Arc::clone(right), + Arc::clone(left), + self.filter().map(JoinFilter::swap), + &self.join_type().swap(), + swap_join_projection( + left.schema().fields().len(), + right.schema().fields().len(), + self.projection.as_deref(), + self.join_type(), + ), + )?; + + // For Semi/Anti joins, swap result will produce same output schema, + // no need to wrap them into additional projection + let plan: Arc<dyn ExecutionPlan> = if matches!( + self.join_type(), + JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ) || self.projection.is_some() + { + Arc::new(new_join) + } else { + reorder_output_after_swap( + Arc::new(new_join), + &self.left().schema(), + &self.right().schema(), + )? + }; + + Ok(plan) + } +} + +impl DisplayAs for NestedLoopJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let display_filter = self.filter.as_ref().map_or_else( + || "".to_string(), + |f| format!(", filter={}", f.expression()), + ); + let display_projections = if self.contains_projection() { + format!( + ", projection=[{}]", + self.projection + .as_ref() + .unwrap() + .iter() + .map(|index| format!( + "{}@{}", + self.join_schema.fields().get(*index).unwrap().name(), + index + )) + .collect::<Vec<_>>() + .join(", ") + ) + } else { + "".to_string() + }; + write!( + f, + "NestedLoopJoinExec: join_type={:?}{}{}", + self.join_type, display_filter, display_projections + ) + } + DisplayFormatType::TreeRender => { + if *self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type) + } else { + Ok(()) + } + } + } + } +} + +impl ExecutionPlan for NestedLoopJoinExec { + fn name(&self) -> &'static str { + "NestedLoopJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::UnspecifiedDistribution, + ]) + } + + fn maintains_input_order(&self) -> Vec<bool> { + Self::maintains_input_order(self.join_type) + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + // Apply to join filter expressions if present + crate::apply_expression_roots( + self.filter.iter().map(|filter| filter.expression()), + f, + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + build_side_data: Default::default(), + left_spill_data: Arc::new(OnceAsync::default()), + cache: Arc::clone(&self.cache), + filter: self.filter.clone(), + join_type: self.join_type, + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + })) + } + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + NestedLoopJoinExecBuilder::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.join_type, + ) + .with_filter(self.filter.clone()) + .with_projection_ref(self.projection.clone()) + .build()?, + )), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + assert_eq_or_internal_err!( + self.left.output_partitioning().partition_count(), + 1, + "Invalid NestedLoopJoinExec, the output partition count of the left child must be 1,\ + consider using CoalescePartitionsExec or the EnforceDistribution rule" + ); + + let metrics = NestedLoopJoinMetrics::new(&self.metrics, partition); + let batch_size = context.session_config().batch_size(); + + // update column indices to reflect the projection + let column_indices_after_projection = match self.projection.as_ref() { + Some(projection) => projection + .iter() + .map(|i| self.column_indices[*i].clone()) + .collect(), + None => self.column_indices.clone(), + }; + + let right_partition_count = self.right().output_partitioning().partition_count(); + + // Always try to buffer all left data in memory via OnceFut. + // If that fails with OOM, the stream will fallback to memory-limited + // mode (if conditions allow). + let load_reservation = + MemoryConsumer::new(format!("NestedLoopJoinLoad[{partition}]")) + .register(context.memory_pool()); + + let build_side_data = self.build_side_data.try_once(|| { + let stream = self.left.execute(0, Arc::clone(&context))?; + + Ok(collect_left_input( + stream, + metrics.join_metrics.clone(), + load_reservation, + need_produce_result_in_final(self.join_type), + right_partition_count, + )) + })?; + + let probe_side_data = self.right.execute(partition, Arc::clone(&context))?; + + // Determine if OOM fallback to memory-limited mode is possible. + // Conditions: + // 1. Disk manager supports temp files (needed for spilling). + // 2. FULL join with multiple right partitions is not yet supported + // in the fallback path. FULL join needs to track BOTH left-side + // matches (for unmatched left rows) AND right-side matches (for + // unmatched right rows). The fallback path builds a per-partition + // `JoinLeftData` with `probe_threads_counter == 1`, so each + // partition emits unmatched left rows based only on its own + // right-side matches, producing incorrect duplicate output for + // left rows that match in another partition. Other join types + // that need only one-sided final emission (LEFT, LEFT SEMI, + // LEFT ANTI, LEFT MARK) have a similar latent issue in the + // fallback path which predates this change; tracking is out of + // scope for this PR. + let full_join_multi_partition = + matches!(self.join_type, JoinType::Full) && right_partition_count > 1; + let spill_state = if context.runtime_env().disk_manager.tmp_files_enabled() + && !full_join_multi_partition + { + SpillState::Pending { + left_plan: Arc::clone(&self.left), + task_context: Arc::clone(&context), + left_spill_data: Arc::clone(&self.left_spill_data), + } + } else { + SpillState::Disabled + }; + + Ok(Box::pin(NestedLoopJoinStream::new( + self.schema(), + self.filter.clone(), + self.join_type, + probe_side_data, + build_side_data, + column_indices_after_projection, + metrics, + batch_size, + spill_state, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + // Left side is always broadcast, so it always needs overall stats. + // Right side is partitioned, so it needs per-partition stats. + vec![ChildStats::At(None), ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + // NestedLoopJoinExec is designed for joins without equijoin keys in the + // ON clause (e.g., `t1 JOIN t2 ON (t1.v1 + t2.v1) % 2 = 0`). Any join + // predicates are stored in `self.filter`, but `estimate_join_statistics` + // currently doesn't support selectivity estimation for such arbitrary + // filter expressions. We pass an empty join column list, which means + // the cardinality estimation cannot use column statistics and returns + // unknown row counts. + let join_columns = Vec::new(); + + let left_stats = input_stats[0].as_ref().clone(); + let right_stats = input_stats[1].as_ref().clone(); + + let stats = estimate_join_statistics( + left_stats, + right_stats, + &join_columns, + NullEquality::NullEqualsNothing, + &self.join_type, + &self.join_schema, + )?; + + Ok(Arc::new(stats.project(self.projection.as_ref()))) + } + + /// Tries to push `projection` down through `nested_loop_join`. If possible, performs the + /// pushdown and returns a new [`NestedLoopJoinExec`] as the top plan which has projections + /// as its children. Otherwise, returns `None`. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // TODO: currently if there is projection in NestedLoopJoinExec, we can't push down projection to left or right input. Maybe we can pushdown the mixed projection later. + if self.contains_projection() { + return Ok(None); + } + + let schema = self.schema(); + if let Some(JoinData { + projected_left_child, + projected_right_child, + join_filter, + .. + }) = try_pushdown_through_join_with_column_indices( + projection, + self.left(), + self.right(), + &[], + &schema, + self.filter(), + self.column_indices.as_slice(), + )? { + Ok(Some(Arc::new(NestedLoopJoinExec::try_new( + Arc::new(projected_left_child), + Arc::new(projected_right_child), + join_filter, + self.join_type(), + // Returned early if projection is not None + None, + )?))) + } else { + try_embed_projection(projection, self) + } + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let join_type = crate::joins::proto::join_type_to_proto(*self.join_type()); + + let filter = self + .filter() + .map(|f| crate::joins::proto::join_filter_to_proto(f, ctx)) + .transpose()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin(Box::new( + protobuf::NestedLoopJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + join_type: join_type.into(), + filter, + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl NestedLoopJoinExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin, + "NestedLoopJoinExec", + ); + + let left = ctx.decode_required_child( + join.left.as_deref(), + "NestedLoopJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + join.right.as_deref(), + "NestedLoopJoinExec", + "right", + )?; + + let join_type = crate::joins::proto::join_type_from_proto( + join.join_type, + "NestedLoopJoinExec", + )?; + + let filter = join + .filter + .as_ref() + .map(|f| { + crate::joins::proto::join_filter_from_proto(f, ctx, "NestedLoopJoinExec") + }) + .transpose()?; + + let projection = match join.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + Ok(Arc::new(NestedLoopJoinExec::try_new( + left, right, filter, &join_type, projection, + )?)) + } +} + +impl EmbeddedProjection for NestedLoopJoinExec { + fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self> { + self.with_projection(projection) + } +} + +/// Left (build-side) data +pub(crate) struct JoinLeftData { + /// Build-side data collected to single batch + batch: RecordBatch, + /// Shared bitmap builder for visited left indices + bitmap: SharedBitmapBuilder, + /// Counter of running probe-threads, potentially able to update `bitmap` + probe_threads_counter: AtomicUsize, + /// Memory reservation for tracking batch and bitmap + /// Cleared on `JoinLeftData` drop + /// reservation is cleared on Drop + #[expect(dead_code)] + reservation: MemoryReservation, +} + +impl JoinLeftData { + pub(crate) fn new( + batch: RecordBatch, + bitmap: SharedBitmapBuilder, + probe_threads_counter: AtomicUsize, + reservation: MemoryReservation, + ) -> Self { + Self { + batch, + bitmap, + probe_threads_counter, + reservation, + } + } + + pub(crate) fn batch(&self) -> &RecordBatch { + &self.batch + } + + pub(crate) fn bitmap(&self) -> &SharedBitmapBuilder { + &self.bitmap + } + + /// Decrements counter of running threads, and returns `true` + /// if caller is the last running thread + pub(crate) fn report_probe_completed(&self) -> bool { + self.probe_threads_counter.fetch_sub(1, Ordering::Relaxed) == 1 + } +} + +/// Asynchronously collect input into a single batch, and creates `JoinLeftData` from it +async fn collect_left_input( + stream: SendableRecordBatchStream, + join_metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + with_visited_left_side: bool, + probe_threads_count: usize, +) -> Result<JoinLeftData> { + let schema = stream.schema(); + + // Load all batches and count the rows + let (batches, metrics, reservation) = stream + .try_fold( + (Vec::new(), join_metrics, reservation), + |(mut batches, metrics, reservation), batch| async { + let batch_size = batch.get_array_memory_size(); + // Reserve memory for incoming batch + reservation.try_grow(batch_size)?; + // Update metrics + metrics.build_mem_used.add(batch_size); + metrics.build_input_batches.add(1); + metrics.build_input_rows.add(batch.num_rows()); + // Push batch to output + batches.push(batch); + Ok((batches, metrics, reservation)) + }, + ) + .await?; + + let merged_batch = concat_batches(&schema, &batches)?; + + // Reserve memory for visited_left_side bitmap if required by join type + let visited_left_side = if with_visited_left_side { + let n_rows = merged_batch.num_rows(); + let buffer_size = n_rows.div_ceil(8); + reservation.try_grow(buffer_size)?; + metrics.build_mem_used.add(buffer_size); + + let mut buffer = BooleanBufferBuilder::new(n_rows); + buffer.append_n(n_rows, false); + buffer + } else { + BooleanBufferBuilder::new(0) + }; + + Ok(JoinLeftData::new( + merged_batch, + Mutex::new(visited_left_side), + AtomicUsize::new(probe_threads_count), + reservation, + )) +} + +/// States for join processing. See `poll_next()` comment for more details about +/// state transitions. +#[derive(Debug, Clone, Copy)] +enum NLJState { + BufferingLeft, + FetchingRight, + ProbeRight, + EmitRightUnmatched, + /// Entered exactly once per left chunk, when the probe (right) side is + /// exhausted and probing for the current chunk is finished. This state + /// owns the single [`JoinLeftData::report_probe_completed`] call that + /// decrements the shared probe-threads counter, and records in + /// `is_unmatched_left_emitter` whether this stream is the one responsible + /// for emitting unmatched-left rows. Splitting this decision out of + /// `EmitLeftUnmatched` makes "decrement exactly once" a structural + /// property of the state graph, so the (re-enterable) emit state no longer + /// has to guard against decrementing twice. + ProbeEnd, + EmitLeftUnmatched, + /// Emit unmatched right rows using the global bitmap accumulated across + /// all left chunks. Only used in memory-limited mode for join types that + /// require tracking right-side matches in the final output (RIGHT, FULL, + /// RIGHT SEMI, RIGHT ANTI, RIGHT MARK). + EmitGlobalRightUnmatched, + Done, +} +/// Shared data for the left-side spill fallback. +/// +/// When the in-memory `OnceFut` path fails with OOM, the first partition +/// spills the entire left side to disk. This struct holds the spill file +/// reference so other partitions can read from the same file. +pub(crate) struct LeftSpillData { + /// SpillManager used to read the spill file (has the left schema) + spill_manager: SpillManager, + /// The spill file containing all left-side batches + spill_file: Arc<dyn SpillFile>, + /// Left-side schema + schema: SchemaRef, +} + +/// Tracks the state of the memory-limited spill fallback for NLJ. +/// +/// The NLJ always starts with the standard OnceFut path. If the in-memory +/// load fails with OOM and conditions allow, the operator falls back to a +/// multi-pass strategy where left data is loaded in chunks and the right +/// side is spilled to disk. +pub(crate) enum SpillState { + /// Fallback is not possible (e.g., join type requires global right bitmap, + /// or disk manager is disabled). OOM errors will propagate as-is. + Disabled, + + /// Fallback is possible but not yet triggered. The operator is still + /// attempting the standard OnceFut path. Holds the context needed to + /// initiate fallback if OOM occurs. + Pending { + /// Left child plan for re-execution + left_plan: Arc<dyn ExecutionPlan>, + /// TaskContext for re-execution and SpillManager creation + task_context: Arc<TaskContext>, + /// Shared OnceAsync for left-side spill data. The first partition + /// to initiate fallback spills the left side; others share the file. + left_spill_data: Arc<OnceAsync<LeftSpillData>>, + }, + + /// Fallback has been triggered. Left data is being loaded in chunks + /// and the right side is spilled to disk for re-scanning. + Active(Box<SpillStateActive>), +} + +/// State for active memory-limited spill execution. +/// Boxed inside [`SpillState::Active`] to reduce enum size. +pub(crate) struct SpillStateActive { + /// Shared future for left-side spill data. All partitions wait on + /// the same future — the first to poll triggers the actual spill. + left_spill_fut: OnceFut<LeftSpillData>, + /// Left input stream for incremental chunk reading (from spill file). + /// None until `left_spill_fut` resolves. + left_stream: Option<SendableRecordBatchStream>, + /// Left-side schema (set once `left_spill_fut` resolves) + left_schema: Option<SchemaRef>, + /// Memory reservation for left-side buffering + reservation: MemoryReservation, + /// Accumulated left batches for the current chunk + pending_batches: Vec<RecordBatch>, + /// Right input that spills on the first pass and replays from spill later. + right_input: ReplayableStreamSource, + /// Per-batch accumulated right bitmaps across all left chunks. + /// Index = right batch sequence number (0-based, non-empty batches only). + /// Only populated when `should_track_unmatched_right` is true. + global_right_bitmaps: Vec<BooleanBuffer>, + /// Separate reservation for `global_right_bitmaps`. These buffers live + /// for the full operator lifetime (not per-chunk), so they must be + /// tracked separately from `reservation`, which gets `resize(0)`-ed + /// between chunks. + global_right_bitmaps_reservation: MemoryReservation, + /// Current right batch sequence index within the current pass. + right_batch_index: usize, +} + +impl SpillStateActive { + /// Merge a per-pass right bitmap into the global accumulator at the + /// given batch index, growing the dedicated reservation when seeing + /// a batch index for the first time. + /// + /// On first encounter of `idx`, the bitmap is stored as-is and its + /// size is reserved. On subsequent encounters (later left chunk + /// passes over the same right batch), the existing entry is OR-merged + /// with `values`. Because `bitor` produces a buffer of the same bit + /// length, the reservation does not need to be adjusted on merge. + fn merge_current_right_bitmap(&mut self, idx: usize, values: BooleanBuffer) { + if idx >= self.global_right_bitmaps.len() { + // First encounter of this right batch — account memory and store. + // The bitmap has one bit per right row, so for very large right + // inputs the accumulated size can be non-negligible (e.g., + // 1M rows ≈ 125 KB per batch). + // Use infallible `grow` because we must accept the bitmap to + // preserve correctness — the fallback path has no other recourse. + let bytes = values.len().div_ceil(8); + self.global_right_bitmaps_reservation.grow(bytes); + self.global_right_bitmaps.push(values); + } else { + // Subsequent left chunk pass — OR merge. Same bit length, so + // no reservation adjustment is needed. + self.global_right_bitmaps[idx] = + self.global_right_bitmaps[idx].bitor(&values); + } + } +} + +pub(crate) struct NestedLoopJoinStream { + // ======================================================================== + // PROPERTIES: + // Operator's properties that remain constant + // + // Note: The implementation uses the terms left/build-side table and + // right/probe-side table interchangeably. Treating the left side as the + // build side is a convention in DataFusion: the planner always tries to + // swap the smaller table to the left side. + // ======================================================================== + /// Output schema + pub(crate) output_schema: Arc<Schema>, + /// join filter + pub(crate) join_filter: Option<JoinFilter>, + /// type of the join + pub(crate) join_type: JoinType, + /// the probe-side(right) table data of the nested loop join + /// `Option` is used because memory-limited path requires resetting it. + pub(crate) right_data: Option<SendableRecordBatchStream>, + /// the build-side table data of the nested loop join + pub(crate) left_data: OnceFut<JoinLeftData>, + /// Projection to construct the output schema from the left and right tables. + /// Example: + /// - output_schema: ['a', 'c'] + /// - left_schema: ['a', 'b'] + /// - right_schema: ['c'] + /// + /// The column indices would be [(left, 0), (right, 0)] -- taking the left + /// 0th column and right 0th column can construct the output schema. + /// + /// Note there are other columns ('b' in the example) still kept after + /// projection pushdown; this is because they might be used to evaluate + /// the join filter (e.g., `JOIN ON (b+c)>0`). + pub(crate) column_indices: Vec<ColumnIndex>, + /// Join execution metrics + pub(crate) metrics: NestedLoopJoinMetrics, + + /// `batch_size` from configuration + batch_size: usize, + + /// See comments in [`need_produce_right_in_final`] for more detail + should_track_unmatched_right: bool, + + // ======================================================================== + // STATE FLAGS/BUFFERS: + // Fields that hold intermediate data/flags during execution + // ======================================================================== + /// State Tracking + state: NLJState, + /// Output buffer holds the join result to output. It will emit eagerly when + /// the threshold is reached. + output_buffer: Box<BatchCoalescer>, + /// See comments in [`NLJState::Done`] for its purpose + handled_empty_output: bool, + + // Buffer(left) side + // ----------------- + /// The current buffered left data to join + buffered_left_data: Option<Arc<JoinLeftData>>, + /// Index into the left buffered batch. Used in `ProbeRight` state + left_probe_idx: usize, + /// Index into the left buffered batch. Used in `EmitLeftUnmatched` state + left_emit_idx: usize, + /// Should we go back to `BufferingLeft` state again after `EmitLeftUnmatched` + /// state is over. + left_exhausted: bool, + /// If we can buffer all left data in one pass (false means memory-limited multi-pass) + left_buffered_in_one_pass: bool, + + // Probe(right) side + // ----------------- + /// The current probe batch to process + current_right_batch: Option<RecordBatch>, + // For right join, keep track of matched rows in `current_right_batch` + // Constructed when fetching each new incoming right batch in `FetchingRight` state. + current_right_batch_matched: Option<BooleanArray>, + + /// Memory-limited spill fallback state. See [`SpillState`] for details. + spill_state: SpillState, + + /// Whether this stream is the one responsible for emitting unmatched-left + /// rows for the current left chunk. Set in the [`NLJState::ProbeEnd`] state, + /// which is entered exactly once per chunk and owns the single + /// [`JoinLeftData::report_probe_completed`] call: the stream that drives the + /// shared probe-threads counter to zero (the last to finish probing) becomes + /// the emitter. Because the decrement happens once in `ProbeEnd` rather than + /// in the re-enterable `EmitLeftUnmatched` state, the counter can never be + /// decremented twice, so it cannot reach zero before all partitions finish + /// probing (which would otherwise let a partition emit spurious NULL-padded + /// unmatched-left rows early). + is_unmatched_left_emitter: bool, +} + +pub(crate) struct NestedLoopJoinMetrics { + /// Join execution metrics + pub(crate) join_metrics: BuildProbeJoinMetrics, + /// Selectivity of the join: output_rows / (left_rows * right_rows) + pub(crate) selectivity: RatioMetrics, + /// Spill metrics for memory-limited execution + pub(crate) spill_metrics: SpillMetrics, +} + +impl NestedLoopJoinMetrics { + pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + join_metrics: BuildProbeJoinMetrics::new(partition, metrics), + selectivity: MetricBuilder::new(metrics) + .with_type(MetricType::Summary) + .ratio_metrics("selectivity", partition), + spill_metrics: SpillMetrics::new(metrics, partition), + } + } +} + +impl Stream for NestedLoopJoinStream { + type Item = Result<RecordBatch>; + + /// See the comments [`NestedLoopJoinExec`] for high-level design ideas. + /// + /// # Implementation + /// + /// This function is the entry point of NLJ operator's state machine + /// transitions. The rough state transition graph is as follow, for more + /// details see the comment in each state's matching arm. + /// + /// ============================ + /// State transition graph: + /// ============================ + /// + /// (start) --> BufferingLeft + /// ---------------------------- + /// BufferingLeft → FetchingRight + /// + /// FetchingRight → ProbeRight (if right batch available) + /// FetchingRight → ProbeEnd (if right exhausted) + /// + /// ProbeRight → ProbeRight (next left row or after yielding output) + /// ProbeRight → EmitRightUnmatched (for special join types like right join) + /// ProbeRight → FetchingRight (done with the current right batch) + /// + /// EmitRightUnmatched → FetchingRight + /// + /// ProbeEnd → EmitLeftUnmatched (records whether this stream is the + /// unmatched-left emitter, then always continues to EmitLeftUnmatched) + /// + /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each + /// iteration) + /// EmitLeftUnmatched → Done (if finished) + /// ---------------------------- + /// Done → (end) + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Self::Item>> { + loop { + match self.state { + // # NLJState transitions + // --> FetchingRight + // This state will prepare the left side batches, next state + // `FetchingRight` is responsible for preparing a single probe + // 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(()) => continue, + ControlFlow::Break(poll) => return poll, + } + } + + // # NLJState transitions: + // 1. --> ProbeRight + // Start processing the join for the newly fetched right + // batch. + // 2. --> ProbeEnd: When the right side input is exhausted, + // probing for the current left chunk is finished. + // + // After fetching a new batch from the right side, it will + // process all rows from the buffered left data: + // ```text + // for batch in right_side: + // for row in left_buffer: + // join(batch, row) + // ``` + // Note: the implementation does this step incrementally, + // instead of materializing all intermediate Cartesian products + // at once in memory. + // + // So after the right side input is exhausted, the join phase + // for the current buffered left data is finished. We go to the + // `ProbeEnd` state, which records probe completion before the + // `EmitLeftUnmatched` phase checks if there is any special + // 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(()) => continue, + ControlFlow::Break(poll) => return poll, + } + } + + // NLJState transitions: + // 1. --> ProbeRight(1) + // If we have already buffered enough output to yield, it + // will first give back control to the parent state machine, + // then resume at the same place. + // 2. --> ProbeRight(2) + // After probing one right batch, and evaluating the + // join filter on (left-row x right-batch), it will advance + // to the next left row, then re-enter the current state and + // continue joining. + // 3. --> FetchRight + // After it has done with the current right batch (to join + // with all rows in the left buffer), it will go to + // FetchRight state to check what to do next. + NLJState::ProbeRight => { + 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_probe_right() { + ControlFlow::Continue(()) => continue, + ControlFlow::Break(poll) => { + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + + // In the `current_right_batch_matched` bitmap, all trues mean + // it has been output by the join. In this state we have to + // output unmatched rows for current right batch (with null + // padding for left relation) + // Precondition: we have checked the join type so that it's + // possible to output right unmatched (e.g. it's right join) + NLJState::EmitRightUnmatched => { + 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_emit_right_unmatched() { + ControlFlow::Continue(()) => continue, + ControlFlow::Break(poll) => { + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + + // NLJState transitions: + // 1. --> EmitLeftUnmatched + // Probing for the current left chunk is finished. Report + // probe completion exactly once (decrementing the shared + // probe-threads counter) and record whether this stream is + // the unmatched-left emitter, then always advance to + // `EmitLeftUnmatched`. + NLJState::ProbeEnd => { + 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_probe_end() { + ControlFlow::Continue(()) => continue, + ControlFlow::Break(poll) => { + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + + // NLJState transitions: + // 1. --> EmitLeftUnmatched(1) + // If we have already buffered enough output to yield, it + // will first give back control to the parent state machine, + // then resume at the same place. + // 2. --> EmitLeftUnmatched(2) + // After processing some unmatched rows, it will re-enter + // the same state, to check if there are any more final + // results to output. + // 3. --> Done + // It has processed all data, go to the final state and ready + // to exit. + // 4. --> BufferingLeft (memory-limited mode only) + // When left data was loaded in chunks and more chunks remain, + // go back to BufferingLeft to load the next chunk. + NLJState::EmitLeftUnmatched => { + 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_emit_left_unmatched() { + ControlFlow::Continue(()) => continue, + ControlFlow::Break(poll) => { + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + + // Replay all right batches from spill and emit unmatched + // right rows using the global bitmap accumulated across all + // left chunks. Only entered in memory-limited mode for join + // types where `should_track_unmatched_right` is true + // (RIGHT, FULL, RIGHT SEMI, RIGHT ANTI, RIGHT MARK). + 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(()) => continue, + ControlFlow::Break(poll) => { + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + + // The final state and the exit point + NLJState::Done => { + debug!("[NLJState] Entering: {:?}", self.state); + + // stop on drop + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); + // counting it in join timer due to there might be some + // final resout batches to output in this state + + let poll = self.handle_done(); + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + } +} + +impl RecordBatchStream for NestedLoopJoinStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.output_schema) + } +} + +impl NestedLoopJoinStream { + #[expect(clippy::too_many_arguments)] + pub(crate) fn new( + schema: Arc<Schema>, + filter: Option<JoinFilter>, + join_type: JoinType, + right_data: SendableRecordBatchStream, + left_data: OnceFut<JoinLeftData>, + column_indices: Vec<ColumnIndex>, + metrics: NestedLoopJoinMetrics, + batch_size: usize, + spill_state: SpillState, + ) -> Self { + Self { + output_schema: Arc::clone(&schema), + join_filter: filter, + join_type, + right_data: Some(right_data), + column_indices, + left_data, + metrics, + buffered_left_data: None, + output_buffer: Box::new(BatchCoalescer::new(schema, batch_size)), + batch_size, + current_right_batch: None, + current_right_batch_matched: None, + state: NLJState::BufferingLeft, + left_probe_idx: 0, + left_emit_idx: 0, + left_exhausted: false, + left_buffered_in_one_pass: true, + handled_empty_output: false, + should_track_unmatched_right: need_produce_right_in_final(join_type), + spill_state, + is_unmatched_left_emitter: false, + } + } + + /// Returns true if this stream is operating in memory-limited mode + fn is_memory_limited(&self) -> bool { + matches!(self.spill_state, SpillState::Active(_)) + } + + /// Check if we can fall back to memory-limited mode on this error. + fn can_fallback_to_spill(&self, error: &datafusion_common::DataFusionError) -> bool { + matches!(self.spill_state, SpillState::Pending { .. }) + && matches!( + error.find_root(), + datafusion_common::DataFusionError::ResourcesExhausted(_) + ) + } + + /// Switch from the standard OnceFut path to memory-limited mode. + /// + /// Uses the shared `left_spill_data` OnceAsync so that only the first + /// partition to reach this point re-executes the left child and spills + /// it to disk. Other partitions share the same spill file. + fn initiate_fallback(&mut self) -> Result<()> { + // Take ownership of Pending state + let (left_plan, context, left_spill_data) = + match std::mem::replace(&mut self.spill_state, SpillState::Disabled) { + SpillState::Pending { + left_plan, + task_context, + left_spill_data, + } => (left_plan, task_context, left_spill_data), + _ => { + return internal_err!( + "initiate_fallback called in non-Pending spill state" + ); + } + }; + + // Use OnceAsync to ensure only the first partition spills the left + // side. Other partitions will get the same OnceFut that resolves + // to the shared spill file. + let left_spill_fut = left_spill_data.try_once(|| { + let plan = Arc::clone(&left_plan); + let ctx = Arc::clone(&context); + let spill_metrics = self.metrics.spill_metrics.clone(); + Ok(async move { + let mut stream = plan.execute(0, Arc::clone(&ctx))?; + let schema = stream.schema(); + let left_spill_manager = SpillManager::new( + ctx.runtime_env(), + spill_metrics, + Arc::clone(&schema), + ) + .with_compression_type(ctx.session_config().spill_compression()); + + let result = left_spill_manager + .spill_record_batch_stream_and_return_max_batch_memory( + &mut stream, + "NestedLoopJoin left spill", + ) + .await?; + + match result { + Some((file, _max_batch_memory)) => Ok(LeftSpillData { + spill_manager: left_spill_manager, + spill_file: file, + schema, + }), + None => { + internal_err!("Left side produced no data to spill") + } + } + }) + })?; + + // Create reservation with can_spill for fair memory allocation + let reservation = MemoryConsumer::new("NestedLoopJoinLoad[fallback]".to_string()) + .with_can_spill(true) + .register(context.memory_pool()); + + // Separate reservation for the global right bitmaps. These buffers + // persist across all left chunks, whereas `reservation` is reset + // between chunks via `resize(0)`. + let global_right_bitmaps_reservation = + MemoryConsumer::new("NestedLoopJoinGlobalRightBitmaps".to_string()) + .register(context.memory_pool()); + + // Create SpillManager for right-side spilling + let right_schema = self + .right_data + .as_ref() + .expect("right_data must be present before fallback") + .schema(); + let right_data = self + .right_data + .take() + .expect("right_data must be present before fallback"); + let right_spill_manager = SpillManager::new( + context.runtime_env(), + self.metrics.spill_metrics.clone(), + right_schema, + ) + .with_compression_type(context.session_config().spill_compression()); + + self.spill_state = SpillState::Active(Box::new(SpillStateActive { + left_spill_fut, + left_stream: None, + left_schema: None, + reservation, + pending_batches: Vec::new(), + right_input: ReplayableStreamSource::new( + right_data, + right_spill_manager, + "NestedLoopJoin right spill", + ), + global_right_bitmaps: Vec::new(), + global_right_bitmaps_reservation, + right_batch_index: 0, + })); + + // State stays BufferingLeft — next poll will enter + // handle_buffering_left_memory_limited via is_memory_limited() check + self.state = NLJState::BufferingLeft; + + Ok(()) + } + + // ==== State handler functions ==== + + /// Handle BufferingLeft state - prepare left side batches. + /// + /// In standard mode, uses OnceFut to load all left data at once. + /// In memory-limited mode, incrementally buffers left batches until the + /// memory budget is reached or the left stream is exhausted. + fn handle_buffering_left( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + if self.is_memory_limited() { + self.handle_buffering_left_memory_limited(cx) + } else { + // Standard path: use OnceFut + match self.left_data.get_shared(cx) { + Poll::Ready(Ok(left_data)) => { + self.buffered_left_data = Some(left_data); + self.left_exhausted = true; + self.state = NLJState::FetchingRight; + ControlFlow::Continue(()) + } + Poll::Ready(Err(e)) => { + if self.can_fallback_to_spill(&e) { + debug!( + "NestedLoopJoin: OnceFut failed with OOM, \ + falling back to memory-limited mode" + ); + match self.initiate_fallback() { + Ok(()) => ControlFlow::Continue(()), + Err(fallback_err) => { + ControlFlow::Break(Poll::Ready(Some(Err(fallback_err)))) + } + } + } else { + ControlFlow::Break(Poll::Ready(Some(Err(e)))) + } + } + Poll::Pending => ControlFlow::Break(Poll::Pending), + } + } + } + + /// Memory-limited path for handle_buffering_left. + /// + /// Incrementally polls the left stream and accumulates batches until: + /// - Memory reservation fails (chunk is full, more data remains) + /// - Left stream is exhausted (this is the last/only chunk) + fn handle_buffering_left_memory_limited( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + let SpillState::Active(active) = &mut self.spill_state else { + unreachable!( + "handle_buffering_left_memory_limited called without Active spill state" + ); + }; + + // On first entry (or after re-entry for a new chunk pass when + // left_stream was consumed), wait for the shared left spill + // future to resolve and then open a stream from the spill file. + if active.left_stream.is_none() { + match active.left_spill_fut.get_shared(cx) { + Poll::Ready(Ok(spill_data)) => { + match spill_data + .spill_manager + .read_spill_as_stream(Arc::clone(&spill_data.spill_file), None) + { + Ok(stream) => { + active.left_schema = Some(Arc::clone(&spill_data.schema)); + active.left_stream = Some(stream); + } + Err(e) => { + return ControlFlow::Break(Poll::Ready(Some(Err(e)))); + } + } + } + Poll::Ready(Err(e)) => { + return ControlFlow::Break(Poll::Ready(Some(Err(e)))); + } + Poll::Pending => { + return ControlFlow::Break(Poll::Pending); + } + } + } + + let left_stream = active + .left_stream + .as_mut() + .expect("left_stream must be set after spill future resolves"); + + // Poll left stream for more batches. + // Note: pending_batches may already contain a batch from the + // previous chunk iteration (the batch that triggered the memory limit). + loop { + match left_stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + if batch.num_rows() == 0 { + continue; + } + let batch_rows = batch.num_rows(); + let batch_size = batch.get_array_memory_size(); + let can_grow = active.reservation.try_grow(batch_size).is_ok(); + + if !can_grow && !active.pending_batches.is_empty() { + // Memory limit reached and we already have data. + // Push this batch into pending (it's already in memory) + // and stop buffering for this chunk. + active.pending_batches.push(batch); + self.left_exhausted = false; + self.left_buffered_in_one_pass = false; + break; + } else if !can_grow { + // No pending batches yet — we must accept this batch + // to make progress, even if it exceeds the budget. + active.reservation.grow(batch_size); + } + + self.metrics.join_metrics.build_mem_used.add(batch_size); + self.metrics.join_metrics.build_input_batches.add(1); + self.metrics.join_metrics.build_input_rows.add(batch_rows); + active.pending_batches.push(batch); + } + Poll::Ready(Some(Err(e))) => { + return ControlFlow::Break(Poll::Ready(Some(Err(e)))); + } + Poll::Ready(None) => { + // Left stream exhausted + self.left_exhausted = true; + break; + } + Poll::Pending => { + return ControlFlow::Break(Poll::Pending); + } + } + } + + // 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 { + active.left_stream = None; + } + + if active.pending_batches.is_empty() { + // No data at all — go directly to Done + self.left_exhausted = true; + self.state = NLJState::Done; + return ControlFlow::Continue(()); + } + + let merged_batch = match concat_batches( + active + .left_schema + .as_ref() + .expect("left_schema must be set"), + &active.pending_batches, + ) { + Ok(batch) => batch, + Err(e) => { + return ControlFlow::Break(Poll::Ready(Some(Err(e.into())))); + } + }; + active.pending_batches.clear(); + + // Build visited bitmap if needed for this join type + let with_visited = need_produce_result_in_final(self.join_type); + let n_rows = merged_batch.num_rows(); + let visited_left_side = if with_visited { + let buffer_size = n_rows.div_ceil(8); + // Use infallible grow for bitmap — it's small + active.reservation.grow(buffer_size); + self.metrics.join_metrics.build_mem_used.add(buffer_size); + let mut buffer = BooleanBufferBuilder::new(n_rows); + buffer.append_n(n_rows, false); + buffer + } else { + BooleanBufferBuilder::new(0) + }; + + // Create an empty reservation for JoinLeftData's RAII field. + // The actual memory tracking is managed by the Active state's reservation. + let dummy_reservation = active.reservation.new_empty(); + + let left_data = JoinLeftData::new( + merged_batch, + Mutex::new(visited_left_side), + // In memory-limited mode, only 1 probe thread per chunk + AtomicUsize::new(1), + dummy_reservation, + ); + + self.buffered_left_data = Some(Arc::new(left_data)); + + active.right_batch_index = 0; + match active.right_input.open_pass() { + Ok(stream) => { + self.right_data = Some(stream); + } + Err(e) => { + return ControlFlow::Break(Poll::Ready(Some(Err(e)))); + } + } + + self.state = NLJState::FetchingRight; + ControlFlow::Continue(()) + } + + /// Handle FetchingRight state - fetch next right batch and prepare for processing. + /// + /// In memory-limited mode during the first pass, each right batch is also + /// written to a spill file so it can be re-read on subsequent passes. + fn handle_fetching_right( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + 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(()); + } + + 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)); + } + + 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(()) + } + }, + Poll::Pending => ControlFlow::Break(Poll::Pending), + } + } + + /// Handle ProbeRight state - process current probe batch + fn handle_probe_right(&mut self) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + // Return any completed batches first + if let Some(poll) = self.maybe_flush_ready_batch() { + return ControlFlow::Break(poll); + } + + // Process current probe state + match self.process_probe_batch() { + // State unchanged (ProbeRight) + // Continue probing until we have done joining the + // current right batch with all buffered left rows. + Ok(true) => ControlFlow::Continue(()), + // To next FetchRightState + // We have finished joining + // (cur_right_batch x buffered_left_batches) + Ok(false) => { + // Left exhausted, transition to FetchingRight + self.left_probe_idx = 0; + + // Selectivity Metric: Update total possibilities for the batch (left_rows * right_rows) + // If memory-limited execution is implemented, this logic must be updated accordingly. + if let (Ok(left_data), Some(right_batch)) = + (self.get_left_data(), self.current_right_batch.as_ref()) + { + let left_rows = left_data.batch().num_rows(); + let right_rows = right_batch.num_rows(); + self.metrics.selectivity.add_total(left_rows * right_rows); + } + + if self.should_track_unmatched_right { + debug_assert!( + self.current_right_batch_matched.is_some(), + "If it's required to track matched rows in the right input, the right bitmap must be present" + ); + self.state = NLJState::EmitRightUnmatched; + } else { + self.current_right_batch = None; + self.state = NLJState::FetchingRight; + } + ControlFlow::Continue(()) + } + Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + } + } + + /// Handle EmitRightUnmatched state - emit unmatched right rows. + /// + /// In memory-limited mode, instead of emitting unmatched right rows + /// per-batch (which would be incorrect since more left chunks may + /// match those rows), we merge the bitmap into the global accumulator + /// and defer emission to `EmitGlobalRightUnmatched`. + fn handle_emit_right_unmatched( + &mut self, + ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + // In memory-limited mode, merge bitmap into global and move on + if self.is_memory_limited() { + debug_assert!( + self.current_right_batch_matched.is_some(), + "right bitmap must be present" + ); + let bitmap = std::mem::take(&mut self.current_right_batch_matched) + .expect("right bitmap should be available"); + let (values, _nulls) = bitmap.into_parts(); + + if let SpillState::Active(ref mut active) = self.spill_state { + let idx = active.right_batch_index; + active.merge_current_right_bitmap(idx, values); + active.right_batch_index += 1; + } + + self.current_right_batch = None; + self.state = NLJState::FetchingRight; + return ControlFlow::Continue(()); + } + + // Standard (single-pass) mode: emit unmatched right rows immediately + // Return any completed batches first + if let Some(poll) = self.maybe_flush_ready_batch() { + return ControlFlow::Break(poll); + } + + debug_assert!( + self.current_right_batch_matched.is_some() + && self.current_right_batch.is_some(), + "This state is yielding output for unmatched rows in the current right batch, so both the right batch and the bitmap must be present" + ); + match self.process_right_unmatched() { + Ok(Some(batch)) => match self.output_buffer.push_batch(batch) { + Ok(()) => { + debug_assert!(self.current_right_batch.is_none()); + self.state = NLJState::FetchingRight; + ControlFlow::Continue(()) + } + Err(e) => ControlFlow::Break(Poll::Ready(Some(arrow_err!(e)))), + }, + Ok(None) => { + debug_assert!(self.current_right_batch.is_none()); + self.state = NLJState::FetchingRight; + ControlFlow::Continue(()) + } + Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + } + } + + /// Handle ProbeEnd state - record probe completion for the current chunk. + /// + /// Entered exactly once per left chunk, when the right side is exhausted. + /// This is the single place that decrements the shared probe-threads counter + /// via [`JoinLeftData::report_probe_completed`]: the stream that drives the + /// counter to zero (the last to finish probing) is the one responsible for + /// emitting unmatched-left rows, recorded in `is_unmatched_left_emitter`. + /// + /// Owning the decrement here — rather than in the re-enterable + /// `EmitLeftUnmatched` state — makes "decrement exactly once per stream" a + /// structural property of the state graph, so the counter cannot reach zero + /// before all partitions finish probing (which would let a partition emit + /// spurious NULL-padded unmatched-left rows early). + /// + /// Always transitions to `EmitLeftUnmatched`. + fn handle_probe_end(&mut self) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + // Decrement the shared counter exactly once for this stream/chunk. The + // last stream to finish probing (the one that drives the counter to + // zero) becomes the unmatched-left emitter. + let is_emitter = match self.get_left_data() { + Ok(left_data) => left_data.report_probe_completed(), + Err(e) => return ControlFlow::Break(Poll::Ready(Some(Err(e)))), + }; + self.is_unmatched_left_emitter = is_emitter; + self.state = NLJState::EmitLeftUnmatched; + ControlFlow::Continue(()) + } + + /// Handle EmitLeftUnmatched state - emit unmatched left rows. + /// + /// In memory-limited mode, after processing all unmatched rows for the + /// current left chunk, transitions back to `BufferingLeft` to load the + /// next chunk (if the left stream is not yet exhausted). + fn handle_emit_left_unmatched( + &mut self, + ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + // Return any completed batches first + if let Some(poll) = self.maybe_flush_ready_batch() { + return ControlFlow::Break(poll); + } + + // Process current unmatched state + match self.process_left_unmatched() { + // State unchanged (EmitLeftUnmatched) + // Continue processing until we have processed all unmatched rows + Ok(true) => ControlFlow::Continue(()), + // We have finished processing all unmatched rows for this chunk + Ok(false) => match self.output_buffer.finish_buffered_batch() { + Ok(()) => { + // Flush any completed batch before transitioning. + // This is critical for the memory-limited path: the + // ProbeRight results must be emitted before we discard + // the current chunk and load the next one. + if let Some(poll) = self.maybe_flush_ready_batch() { + return ControlFlow::Break(poll); + } + + if !self.left_exhausted && self.is_memory_limited() { + // More left data to process — free current chunk and + // go back to BufferingLeft for the next chunk + if let SpillState::Active(ref active) = self.spill_state { + active.reservation.resize(0); + } + self.buffered_left_data = None; + self.left_probe_idx = 0; + self.left_emit_idx = 0; + // Each memory-limited chunk gets a fresh per-chunk + // `JoinLeftData`/counter; `is_unmatched_left_emitter` is + // recomputed when `ProbeEnd` is re-entered for the next + // chunk, so it does not need to be reset here. + self.state = NLJState::BufferingLeft; + } else if self.is_memory_limited() + && self.should_track_unmatched_right + { + // All left chunks done — emit global right unmatched. + // Drop the exhausted right stream so that + // EmitGlobalRightUnmatched opens a fresh replay pass + // from the spill file. (process_left_unmatched_range + // already ran with right_data still set, so its + // schema access is not affected.) + self.right_data = None; + self.state = NLJState::EmitGlobalRightUnmatched; + } else { + self.state = NLJState::Done; + } + ControlFlow::Continue(()) + } + Err(e) => ControlFlow::Break(Poll::Ready(Some(arrow_err!(e)))), + }, + Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + } + } + + /// Handle EmitGlobalRightUnmatched state. + /// + /// Replays all right batches from the spill file and emits unmatched + /// right rows using the global bitmap accumulated across all left chunks. + fn handle_emit_global_right_unmatched( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> ControlFlow<Poll<Option<Result<RecordBatch>>>> { + // Flush any completed batches first + if let Some(poll) = self.maybe_flush_ready_batch() { + return ControlFlow::Break(poll); + } + + // On first entry, open a new replay pass on the right input + if self.right_data.is_none() { + let SpillState::Active(ref mut active) = self.spill_state else { + unreachable!("EmitGlobalRightUnmatched without Active spill state"); + }; + active.right_batch_index = 0; + match active.right_input.open_pass() { + Ok(stream) => { + self.right_data = Some(stream); + } + Err(e) => { + return ControlFlow::Break(Poll::Ready(Some(Err(e)))); + } + } + } + + // Poll the replay stream for the next right batch + match self + .right_data + .as_mut() + .expect("right_data must be present") + .poll_next_unpin(cx) + { + Poll::Ready(Some(Ok(right_batch))) => { + if right_batch.num_rows() == 0 { + return ControlFlow::Continue(()); + } + + let SpillState::Active(ref mut active) = self.spill_state else { + unreachable!(); + }; + let idx = active.right_batch_index; + active.right_batch_index += 1; + + // Build BooleanArray from the global bitmap + let bitmap = if idx < active.global_right_bitmaps.len() { + BooleanArray::new(active.global_right_bitmaps[idx].clone(), None) + } else { + // Batch never seen — treat all rows as unmatched + BooleanArray::new( + BooleanBuffer::new_unset(right_batch.num_rows()), + None, + ) + }; + + let left_schema = Arc::clone( + active + .left_schema + .as_ref() + .expect("left_schema must be set"), + ); + + match build_unmatched_batch( + &self.output_schema, + &right_batch, + bitmap, + &left_schema, + &self.column_indices, + self.join_type, + JoinSide::Right, + ) { + Ok(Some(batch)) => match self.output_buffer.push_batch(batch) { + Ok(()) => ControlFlow::Continue(()), + Err(e) => ControlFlow::Break(Poll::Ready(Some(arrow_err!(e)))), + }, + Ok(None) => ControlFlow::Continue(()), + Err(e) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + } + } + Poll::Ready(Some(Err(e))) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), + Poll::Ready(None) => { + // All right batches replayed + match self.output_buffer.finish_buffered_batch() { + Ok(()) => { + self.state = NLJState::Done; + ControlFlow::Continue(()) + } + Err(e) => ControlFlow::Break(Poll::Ready(Some(arrow_err!(e)))), + } + } + Poll::Pending => ControlFlow::Break(Poll::Pending), + } + } + + /// Handle Done state - final state processing + fn handle_done(&mut self) -> Poll<Option<Result<RecordBatch>>> { + // Return any remaining completed batches before final termination + if let Some(poll) = self.maybe_flush_ready_batch() { + return poll; + } + + // HACK for the doc test in https://github.com/apache/datafusion/blob/main/datafusion/core/src/dataframe/mod.rs#L1265 + // If this operator directly return `Poll::Ready(None)` + // for empty result, the final result will become an empty + // batch with empty schema, however the expected result + // should be with the expected schema for this operator + if !self.handled_empty_output { + let zero_count = Count::new(); + if *self.metrics.join_metrics.baseline.output_rows() == zero_count { + let empty_batch = RecordBatch::new_empty(Arc::clone(&self.output_schema)); + self.handled_empty_output = true; + return Poll::Ready(Some(Ok(empty_batch))); + } + } + + Poll::Ready(None) + } + + // ==== Core logic handling for each state ==== + + /// Returns bool to indicate should it continue probing + /// true -> continue in the same ProbeRight state + /// false -> It has done with the (buffered_left x cur_right_batch), go to + /// next state (ProbeRight) + fn process_probe_batch(&mut self) -> Result<bool> { + let left_data = Arc::clone(self.get_left_data()?); + let right_batch = self + .current_right_batch + .as_ref() + .ok_or_else(|| internal_datafusion_err!("Right batch should be available"))? + .clone(); + + // stop probing, the caller will go to the next state + if self.left_probe_idx >= left_data.batch().num_rows() { + return Ok(false); + } + + // ======== + // Join (l_row x right_batch) + // and push the result into output_buffer + // ======== + + // Special case: + // When the right batch is very small, join with multiple left rows at once, + // + // The regular implementation is not efficient if the plan's right child is + // very small (e.g. 1 row total), because inside the inner loop of NLJ, it's + // handling one input right batch at once, if it's not large enough, the + // overheads like filter evaluation can't be amortized through vectorization. + debug_assert_ne!( + right_batch.num_rows(), + 0, + "When fetching the right batch, empty batches will be skipped" + ); + + let l_row_cnt_ratio = self.batch_size / right_batch.num_rows(); + if l_row_cnt_ratio > 10 { + // Calculate max left rows to handle at once. This operator tries to handle + // up to `datafusion.execution.batch_size` rows at once in the intermediate + // batch. + let l_row_count = std::cmp::min( + l_row_cnt_ratio, + left_data.batch().num_rows() - self.left_probe_idx, + ); + + debug_assert!( + l_row_count != 0, + "This function should only be entered when there are remaining left rows to process" + ); + let joined_batch = self.process_left_range_join( + &left_data, + &right_batch, + self.left_probe_idx, + l_row_count, + )?; + + if let Some(batch) = joined_batch { + self.output_buffer.push_batch(batch)?; + } + + self.left_probe_idx += l_row_count; + + return Ok(true); + } + + let l_idx = self.left_probe_idx; + let joined_batch = + self.process_single_left_row_join(&left_data, &right_batch, l_idx)?; + + if let Some(batch) = joined_batch { + self.output_buffer.push_batch(batch)?; + } + + // ==== Prepare for the next iteration ==== + + // Advance left cursor + self.left_probe_idx += 1; + + // Return true to continue probing + Ok(true) + } + + /// Process [l_start_index, l_start_index + l_count) JOIN right_batch + /// Returns a RecordBatch containing the join results (None if empty) + /// + /// Side Effect: If the join type requires, left or right side matched bitmap + /// will be set for matched indices. + fn process_left_range_join( + &mut self, + left_data: &JoinLeftData, + right_batch: &RecordBatch, + l_start_index: usize, + l_row_count: usize, + ) -> Result<Option<RecordBatch>> { + // Construct the Cartesian product between the specified range of left rows + // and the entire right_batch. First, it calculates the index vectors, then + // materializes the intermediate batch, and finally applies the join filter + // to it. + // ----------------------------------------------------------- + let right_rows = right_batch.num_rows(); + let total_rows = l_row_count * right_rows; + + // Build index arrays for cartesian product: left_range X right_batch + let left_indices: UInt32Array = + UInt32Array::from_iter_values((0..l_row_count).flat_map(|i| { + std::iter::repeat_n((l_start_index + i) as u32, right_rows) + })); + let right_indices: UInt32Array = UInt32Array::from_iter_values( + (0..l_row_count).flat_map(|_| 0..right_rows as u32), + ); + + debug_assert!( + left_indices.len() == right_indices.len() + && right_indices.len() == total_rows, + "The length or cartesian product should be (left_size * right_size)", + ); + + // Evaluate the join filter (if any) over an intermediate batch built + // using the filter's own schema/column indices. + let bitmap_combined = if let Some(filter) = &self.join_filter { + // Build the intermediate batch for filter evaluation + let intermediate_batch = if filter.schema.fields().is_empty() { + // Constant predicate (e.g., TRUE/FALSE). Use an empty schema with row_count + create_record_batch_with_empty_schema( + Arc::new((*filter.schema).clone()), + total_rows, + )? + } else { + let mut filter_columns: Vec<Arc<dyn Array>> = + Vec::with_capacity(filter.column_indices().len()); + for column_index in filter.column_indices() { + let array = if column_index.side == JoinSide::Left { + let col = left_data.batch().column(column_index.index); + take(col.as_ref(), &left_indices, None)? + } else { + let col = right_batch.column(column_index.index); + take(col.as_ref(), &right_indices, None)? + }; + filter_columns.push(array); + } + + RecordBatch::try_new(Arc::new((*filter.schema).clone()), filter_columns)? + }; + + let filter_result = filter + .expression() + .evaluate(&intermediate_batch)? + .into_array(intermediate_batch.num_rows())?; + let filter_arr = as_boolean_array(&filter_result)?; + + // Combine with null bitmap to get a unified mask + boolean_mask_from_filter(filter_arr) + } else { + // No filter: all pairs match + BooleanArray::from(vec![true; total_rows]) + }; + + // Update the global left or right bitmap for matched indices + // ----------------------------------------------------------- + + // None means we don't have to update left bitmap for this join type + let mut left_bitmap = if need_produce_result_in_final(self.join_type) { + Some(left_data.bitmap().lock()) + } else { + None + }; + + // 'local' meaning: we want to collect 'is_matched' flag for the current + // right batch, after it has joining all of the left buffer, here it's only + // the partial result for joining given left range + let mut local_right_bitmap = if self.should_track_unmatched_right { + let mut current_right_batch_bitmap = BooleanBufferBuilder::new(right_rows); + // Ensure builder has logical length so set_bit is in-bounds + current_right_batch_bitmap.append_n(right_rows, false); + Some(current_right_batch_bitmap) + } else { + None + }; + + // Set the matched bit for left and right side bitmap + for (i, is_matched) in bitmap_combined.iter().enumerate() { + let is_matched = is_matched.ok_or_else(|| { + internal_datafusion_err!("Must be Some after the previous combining step") + })?; + + let l_index = l_start_index + i / right_rows; + let r_index = i % right_rows; + + if let Some(bitmap) = left_bitmap.as_mut() + && is_matched + { + // Map local index back to absolute left index within the batch + bitmap.set_bit(l_index, true); + } + + if let Some(bitmap) = local_right_bitmap.as_mut() + && is_matched + { + bitmap.set_bit(r_index, true); + } + } + + // Apply the local right bitmap to the global bitmap + if self.should_track_unmatched_right { + // Remember to put it back after update + let global_right_bitmap = + std::mem::take(&mut self.current_right_batch_matched).ok_or_else( + || internal_datafusion_err!("right batch's bitmap should be present"), + )?; + let (buf, nulls) = global_right_bitmap.into_parts(); + debug_assert!(nulls.is_none()); + + let current_right_bitmap = local_right_bitmap + .ok_or_else(|| { + internal_datafusion_err!( + "Should be Some if the current join type requires right bitmap" + ) + })? + .finish(); + let updated_global_right_bitmap = buf.bitor(¤t_right_bitmap); + + self.current_right_batch_matched = + Some(BooleanArray::new(updated_global_right_bitmap, None)); + } + + // For the following join types: only bitmaps are updated; do not emit rows now + if matches!( + self.join_type, + JoinType::LeftAnti + | JoinType::LeftSemi + | JoinType::LeftMark + | JoinType::RightAnti + | JoinType::RightMark + | JoinType::RightSemi + ) { + return Ok(None); + } + + // Build the projected output batch (using output schema/column_indices), + // then apply the bitmap filter to it. + if self.output_schema.fields().is_empty() { + // Empty projection: only row count matters + let row_count = bitmap_combined.true_count(); + return Ok(Some(create_record_batch_with_empty_schema( + Arc::clone(&self.output_schema), + row_count, + )?)); + } + + let mut out_columns: Vec<Arc<dyn Array>> = + Vec::with_capacity(self.output_schema.fields().len()); + for column_index in &self.column_indices { + let array = if column_index.side == JoinSide::Left { + let col = left_data.batch().column(column_index.index); + take(col.as_ref(), &left_indices, None)? + } else { + let col = right_batch.column(column_index.index); + take(col.as_ref(), &right_indices, None)? + }; + out_columns.push(array); + } + let pre_filtered = + RecordBatch::try_new(Arc::clone(&self.output_schema), out_columns)?; + let filtered = filter_record_batch(&pre_filtered, &bitmap_combined)?; + Ok(Some(filtered)) + } + + /// Process a single left row join with the current right batch. + /// Returns a RecordBatch containing the join results (None if empty) + /// + /// Side Effect: If the join type requires, left or right side matched bitmap + /// will be set for matched indices. + fn process_single_left_row_join( + &mut self, + left_data: &JoinLeftData, + right_batch: &RecordBatch, + l_index: usize, + ) -> Result<Option<RecordBatch>> { + let right_row_count = right_batch.num_rows(); + if right_row_count == 0 { + return Ok(None); + } + + let cur_right_bitmap = if let Some(filter) = &self.join_filter { + apply_filter_to_row_join_batch( + left_data.batch(), + l_index, + right_batch, + filter, + )? + } else { + BooleanArray::from(vec![true; right_row_count]) + }; + + self.update_matched_bitmap(l_index, &cur_right_bitmap)?; + + // For the following join types: here we only have to set the left/right + // bitmap, and no need to output result + if matches!( + self.join_type, + JoinType::LeftAnti + | JoinType::LeftSemi + | JoinType::LeftMark + | JoinType::RightAnti + | JoinType::RightMark + | JoinType::RightSemi + ) { + return Ok(None); + } + + if !cur_right_bitmap.has_true() { + // If none of the pairs has passed the join predicate/filter + Ok(None) + } else { + // Use the optimized approach similar to build_intermediate_batch_for_single_left_row + let join_batch = build_row_join_batch( + &self.output_schema, + left_data.batch(), + l_index, + right_batch, + Some(cur_right_bitmap), + &self.column_indices, + JoinSide::Left, + )?; + Ok(join_batch) + } + } + + /// Returns bool to indicate should it continue processing unmatched rows + /// true -> continue in the same EmitLeftUnmatched state + /// false -> next state (Done) + fn process_left_unmatched(&mut self) -> Result<bool> { + let left_data = self.get_left_data()?; + let left_batch = left_data.batch(); + + // ======== + // Check early return conditions + // ======== + + // Early return if join type can't have unmatched rows + let join_type_no_produce_left = !need_produce_result_in_final(self.join_type); + // Stop processing unmatched rows, the caller will go to the next state + let finished = self.left_emit_idx >= left_batch.num_rows(); + + // `ProbeEnd` already recorded whether this stream emits unmatched-left + // rows. Every probe partition passes through this state, but only the + // one that finished probing last is the emitter, so this flag is false + // for the others. + if join_type_no_produce_left || !self.is_unmatched_left_emitter || finished { + return Ok(false); + } + + // ======== + // Process unmatched rows and push the result into output_buffer + // Each time, the number to process is up to batch size + // ======== + let start_idx = self.left_emit_idx; + let end_idx = std::cmp::min(start_idx + self.batch_size, left_batch.num_rows()); + + if let Some(batch) = + self.process_left_unmatched_range(left_data, start_idx, end_idx)? + { + self.output_buffer.push_batch(batch)?; + } + + // ==== Prepare for the next iteration ==== + self.left_emit_idx = end_idx; + + // Return true to continue processing unmatched rows + Ok(true) + } + + /// Process unmatched rows from the left data within the specified range. + /// Returns a RecordBatch containing the unmatched rows (None if empty). + /// + /// # Arguments + /// * `left_data` - The left side data containing the batch and bitmap + /// * `start_idx` - Start index (inclusive) of the range to process + /// * `end_idx` - End index (exclusive) of the range to process + /// + /// # Safety + /// The caller is responsible for ensuring that `start_idx` and `end_idx` are + /// within valid bounds of the left batch. This function does not perform + /// bounds checking. + fn process_left_unmatched_range( + &self, + left_data: &JoinLeftData, + start_idx: usize, + end_idx: usize, + ) -> Result<Option<RecordBatch>> { + if start_idx == end_idx { + return Ok(None); + } + + // Slice both left batch, and bitmap to range [start_idx, end_idx) + // The range is bit index (not byte) + let left_batch = left_data.batch(); + let left_batch_sliced = left_batch.slice(start_idx, end_idx - start_idx); + + // Can this be more efficient? + let mut bitmap_sliced = BooleanBufferBuilder::new(end_idx - start_idx); + bitmap_sliced.append_n(end_idx - start_idx, false); + let bitmap = left_data.bitmap().lock(); + for i in start_idx..end_idx { + assert!( + i - start_idx < bitmap_sliced.capacity(), + "DBG: {start_idx}, {end_idx}" + ); + bitmap_sliced.set_bit(i - start_idx, bitmap.get_bit(i)); + } + let bitmap_sliced = BooleanArray::new(bitmap_sliced.finish(), None); + + let right_schema = self + .right_data + .as_ref() + .expect("right_data must be present when building unmatched batch") + .schema(); + build_unmatched_batch( + &self.output_schema, + &left_batch_sliced, + bitmap_sliced, + &right_schema, + &self.column_indices, + self.join_type, + JoinSide::Left, + ) + } + + /// Process unmatched rows from the current right batch and reset the bitmap. + /// Returns a RecordBatch containing the unmatched right rows (None if empty). + fn process_right_unmatched(&mut self) -> Result<Option<RecordBatch>> { + // ==== Take current right batch and its bitmap ==== + let right_batch_bitmap: BooleanArray = + std::mem::take(&mut self.current_right_batch_matched).ok_or_else(|| { + internal_datafusion_err!("right bitmap should be available") + })?; + + let right_batch = self.current_right_batch.take(); + let cur_right_batch = unwrap_or_internal_err!(right_batch); + + let left_data = self.get_left_data()?; + let left_schema = left_data.batch().schema(); + + let res = build_unmatched_batch( + &self.output_schema, + &cur_right_batch, + right_batch_bitmap, + &left_schema, + &self.column_indices, + self.join_type, + JoinSide::Right, + ); + + // ==== Clean-up ==== + self.current_right_batch_matched = None; + + res + } + + // ==== Utilities ==== + + /// Get the build-side data of the left input, errors if it's None + fn get_left_data(&self) -> Result<&Arc<JoinLeftData>> { + self.buffered_left_data + .as_ref() + .ok_or_else(|| internal_datafusion_err!("LeftData should be available")) + } + + /// Flush the `output_buffer` if there are batches ready to output + /// None if no result batch ready. + fn maybe_flush_ready_batch(&mut self) -> Option<Poll<Option<Result<RecordBatch>>>> { + if self.output_buffer.has_completed_batch() + && let Some(batch) = self.output_buffer.next_completed_batch() + { + // Update output rows for selectivity metric + let output_rows = batch.num_rows(); + self.metrics.selectivity.add_part(output_rows); + + return Some(Poll::Ready(Some(Ok(batch)))); + } + + None + } + + /// After joining (l_index@left_buffer x current_right_batch), it will result + /// in a bitmap (the same length as current_right_batch) as the join match + /// result. Use this bitmap to update the global bitmap, for special join + /// types like full joins. + /// + /// Example: + /// After joining l_index=1 (1-indexed row in the left buffer), and the + /// current right batch with 3 elements, this function will be called with + /// arguments: l_index = 1, r_matched = [false, false, true] + /// - If the join type is FullJoin, the 1-index in the left bitmap will be + /// set to true, and also the right bitmap will be bitwise-ORed with the + /// input r_matched bitmap. + /// - For join types that don't require output unmatched rows, this + /// function can be a no-op. For inner joins, this function is a no-op; for left + /// joins, only the left bitmap may be updated. + fn update_matched_bitmap( + &mut self, + l_index: usize, + r_matched_bitmap: &BooleanArray, + ) -> Result<()> { + let left_data = self.get_left_data()?; + + // 1. Maybe update the left bitmap + if need_produce_result_in_final(self.join_type) && r_matched_bitmap.has_true() { + let mut bitmap = left_data.bitmap().lock(); + bitmap.set_bit(l_index, true); + } + + // 2. Maybe update the right bitmap + if self.should_track_unmatched_right { + debug_assert!(self.current_right_batch_matched.is_some()); + // after bit-wise or, it will be put back + let right_bitmap = std::mem::take(&mut self.current_right_batch_matched) + .ok_or_else(|| { + internal_datafusion_err!("right batch's bitmap should be present") + })?; + let (buf, nulls) = right_bitmap.into_parts(); + debug_assert!(nulls.is_none()); + let updated_right_bitmap = buf.bitor(r_matched_bitmap.values()); + + self.current_right_batch_matched = + Some(BooleanArray::new(updated_right_bitmap, None)); + } + + Ok(()) + } +} + +// ==== Utilities ==== + +/// Apply the join filter between: +/// (l_index th row in left buffer) x (right batch) +/// Returns a bitmap, with successfully joined indices set to true +fn apply_filter_to_row_join_batch( + left_batch: &RecordBatch, + l_index: usize, + right_batch: &RecordBatch, + filter: &JoinFilter, +) -> Result<BooleanArray> { + debug_assert!(left_batch.num_rows() != 0 && right_batch.num_rows() != 0); + + let intermediate_batch = if filter.schema.fields().is_empty() { + // If filter is constant (e.g. literal `true`), empty batch can be used + // in the later filter step. + create_record_batch_with_empty_schema( + Arc::new((*filter.schema).clone()), + right_batch.num_rows(), + )? + } else { + build_row_join_batch( + &filter.schema, + left_batch, + l_index, + right_batch, + None, + &filter.column_indices, + JoinSide::Left, + )? + .ok_or_else(|| internal_datafusion_err!("This function assume input batch is not empty, so the intermediate batch can't be empty too"))? + }; + + let filter_result = filter + .expression() + .evaluate(&intermediate_batch)? + .into_array(intermediate_batch.num_rows())?; + let filter_arr = as_boolean_array(&filter_result)?; + + // Convert boolean array with potential nulls into a unified mask bitmap + let bitmap_combined = boolean_mask_from_filter(filter_arr); + + Ok(bitmap_combined) +} + +/// Convert a boolean filter array into a unified mask bitmap. +/// +/// Caution: The filter result is NOT a bitmap; it contains true/false/null values. +/// For example, `1 < NULL` evaluates to NULL. Therefore, we must combine (AND) +/// the boolean array with its null bitmap to construct a unified bitmap. +#[inline] +fn boolean_mask_from_filter(filter_arr: &BooleanArray) -> BooleanArray { + let (values, nulls) = filter_arr.clone().into_parts(); + match nulls { + Some(nulls) => BooleanArray::new(nulls.inner() & &values, None), + None => BooleanArray::new(values, None), + } +} + +/// This function performs the following steps: +/// 1. Apply filter to probe-side batch +/// 2. Broadcast the left row (build_side_batch\[build_side_index\]) to the +/// filtered probe-side batch +/// 3. Concat them together according to `col_indices`, and return the result +/// (None if the result is empty) +/// +/// Example: +/// build_side_batch: +/// a +/// ---- +/// 1 +/// 2 +/// 3 +/// +/// # 0 index element in the build_side_batch (that is `1`) will be used +/// build_side_index: 0 +/// +/// probe_side_batch: +/// b +/// ---- +/// 10 +/// 20 +/// 30 +/// 40 +/// +/// # After applying it, only index 1 and 3 elements in probe_side_batch will be +/// # kept +/// probe_side_filter: +/// false +/// true +/// false +/// true +/// +/// +/// # Projections to the build/probe side batch, to construct the output batch +/// col_indices: +/// [(left, 0), (right, 0)] +/// +/// build_side: left +/// +/// ==== +/// Result batch: +/// a b +/// ---- +/// 1 20 +/// 1 40 +fn build_row_join_batch( + output_schema: &Schema, + build_side_batch: &RecordBatch, + build_side_index: usize, + probe_side_batch: &RecordBatch, + probe_side_filter: Option<BooleanArray>, + // See [`NLJStream`] struct's `column_indices` field for more detail + col_indices: &[ColumnIndex], + // If the build side is left or right, used to interpret the side information + // in `col_indices` + build_side: JoinSide, +) -> Result<Option<RecordBatch>> { + debug_assert!(build_side != JoinSide::None); + + // TODO(perf): since the output might be projection of right batch, this + // filtering step is more efficient to be done inside the column_index loop + let filtered_probe_batch = if let Some(filter) = probe_side_filter { + &filter_record_batch(probe_side_batch, &filter)? + } else { + probe_side_batch + }; + + if filtered_probe_batch.num_rows() == 0 { + return Ok(None); + } + + // Edge case: downstream operator does not require any columns from this NLJ, + // so allow an empty projection. + // Example: + // SELECT DISTINCT 32 AS col2 + // FROM tab0 AS cor0 + // LEFT OUTER JOIN tab2 AS cor1 + // ON ( NULL ) IS NULL; + if output_schema.fields.is_empty() { + return Ok(Some(create_record_batch_with_empty_schema( + Arc::new(output_schema.clone()), + filtered_probe_batch.num_rows(), + )?)); + } + + let mut columns: Vec<Arc<dyn Array>> = + Vec::with_capacity(output_schema.fields().len()); + + for column_index in col_indices { + let array = if column_index.side == build_side { + // Broadcast the single build-side row to match the filtered + // probe-side batch length + let original_left_array = build_side_batch.column(column_index.index); + + // Use `arrow::compute::take` directly for `List(Utf8View)` rather + // than going through `ScalarValue::to_array_of_size()`, which + // avoids some intermediate allocations. + // + // In other cases, `to_array_of_size()` is faster. + match original_left_array.data_type() { + DataType::List(field) | DataType::LargeList(field) + if field.data_type() == &DataType::Utf8View => + { + let indices_iter = std::iter::repeat_n( + build_side_index as u64, + filtered_probe_batch.num_rows(), + ); + let indices_array = UInt64Array::from_iter_values(indices_iter); + take(original_left_array.as_ref(), &indices_array, None)? + } + _ => { + let scalar_value = ScalarValue::try_from_array( + original_left_array.as_ref(), + build_side_index, + )?; + scalar_value.to_array_of_size(filtered_probe_batch.num_rows())? + } + } + } else { + // Take the filtered probe-side column using compute::take + Arc::clone(filtered_probe_batch.column(column_index.index)) + }; + + columns.push(array); + } + + Ok(Some(RecordBatch::try_new( + Arc::new(output_schema.clone()), + columns, + )?)) +} + +/// Special case for `PlaceHolderRowExec` +/// Minimal example: SELECT 1 WHERE EXISTS (SELECT 1); +// +/// # Return +/// If Some, that's the result batch +/// If None, it's not for this special case. Continue execution. +fn build_unmatched_batch_empty_schema( + output_schema: &SchemaRef, + batch_bitmap: &BooleanArray, + // For left/right/full joins, it needs to fill nulls for another side + join_type: JoinType, +) -> Result<Option<RecordBatch>> { + let result_size = match join_type { + JoinType::Left + | JoinType::Right + | JoinType::Full + | JoinType::LeftAnti + | JoinType::RightAnti => batch_bitmap.false_count(), + JoinType::LeftSemi | JoinType::RightSemi => batch_bitmap.true_count(), + JoinType::LeftMark | JoinType::RightMark => batch_bitmap.len(), + _ => unreachable!(), + }; + + if output_schema.fields().is_empty() { + Ok(Some(create_record_batch_with_empty_schema( + Arc::clone(output_schema), + result_size, + )?)) + } else { + Ok(None) + } +} + +/// Creates an empty RecordBatch with a specific row count. +/// This is useful for cases where we need a batch with the correct schema and row count +/// but no actual data columns (e.g., for constant filters). +fn create_record_batch_with_empty_schema( + schema: SchemaRef, + row_count: usize, +) -> Result<RecordBatch> { + let options = RecordBatchOptions::new() + .with_match_field_names(true) + .with_row_count(Some(row_count)); + + RecordBatch::try_new_with_options(schema, vec![], &options).map_err(|e| { + internal_datafusion_err!("Failed to create empty record batch: {}", e) + }) +} + +/// # Example: +/// batch: +/// a +/// ---- +/// 1 +/// 2 +/// 3 +/// +/// batch_bitmap: +/// ---- +/// false +/// true +/// false +/// +/// another_side_schema: +/// [(b, bool), (c, int32)] +/// +/// join_type: JoinType::Left +/// +/// col_indices: ...(please refer to the comment in `NLJStream::column_indices``) +/// +/// batch_side: right +/// +/// # Walkthrough: +/// +/// This executor is performing a right join, and the currently processed right +/// batch is as above. After joining it with all buffered left rows, the joined +/// entries are marked by the `batch_bitmap`. +/// This method will keep the unmatched indices on the batch side (right), and pad +/// the left side with nulls. The result would be: +/// +/// b c a +/// ------------------------ +/// Null(bool) Null(Int32) 1 +/// Null(bool) Null(Int32) 3 +fn build_unmatched_batch( + output_schema: &SchemaRef, + batch: &RecordBatch, + batch_bitmap: BooleanArray, + // For left/right/full joins, it needs to fill nulls for another side + another_side_schema: &SchemaRef, + col_indices: &[ColumnIndex], + join_type: JoinType, + batch_side: JoinSide, +) -> Result<Option<RecordBatch>> { + // Should not call it for inner joins + debug_assert_ne!(join_type, JoinType::Inner); + debug_assert_ne!(batch_side, JoinSide::None); + + // Handle special case (see function comment) + if let Some(batch) = + build_unmatched_batch_empty_schema(output_schema, &batch_bitmap, join_type)? + { + return Ok(Some(batch)); + } + + match join_type { + JoinType::Full | JoinType::Right | JoinType::Left => { + if join_type == JoinType::Right { + debug_assert_eq!(batch_side, JoinSide::Right); + } + if join_type == JoinType::Left { + debug_assert_eq!(batch_side, JoinSide::Left); + } + + // 1. Filter the batch with *flipped* bitmap + // 2. Fill left side with nulls + let flipped_bitmap = not(&batch_bitmap)?; + + // create a record batch, with left_schema, of only one row of all nulls + let left_null_columns: Vec<Arc<dyn Array>> = another_side_schema + .fields() + .iter() + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + + // Hack: If the left schema is not nullable, the full join result + // might contain null, this is only a temporary batch to construct + // such full join result. + let nullable_left_schema = Arc::new(Schema::new( + another_side_schema + .fields() + .iter() + .map(|field| (**field).clone().with_nullable(true)) + .collect::<Vec<_>>(), + )); + let left_null_batch = if nullable_left_schema.fields.is_empty() { + // Left input can be an empty relation, in this case left relation + // won't be used to construct the result batch (i.e. not in `col_indices`) + create_record_batch_with_empty_schema(nullable_left_schema, 0)? + } else { + RecordBatch::try_new(nullable_left_schema, left_null_columns)? + }; + + debug_assert_ne!(batch_side, JoinSide::None); + let opposite_side = batch_side.negate(); + + build_row_join_batch( + output_schema, + &left_null_batch, + 0, + batch, + Some(flipped_bitmap), + col_indices, + opposite_side, + ) + } + JoinType::RightSemi + | JoinType::RightAnti + | JoinType::LeftSemi + | JoinType::LeftAnti => { + if matches!(join_type, JoinType::RightSemi | JoinType::RightAnti) { + debug_assert_eq!(batch_side, JoinSide::Right); + } + if matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti) { + debug_assert_eq!(batch_side, JoinSide::Left); + } + + let bitmap = if matches!(join_type, JoinType::LeftSemi | JoinType::RightSemi) + { + batch_bitmap.clone() + } else { + not(&batch_bitmap)? + }; + + if !bitmap.has_true() { + return Ok(None); + } + + let mut columns: Vec<Arc<dyn Array>> = + Vec::with_capacity(output_schema.fields().len()); + + for column_index in col_indices { + debug_assert!(column_index.side == batch_side); + + let col = batch.column(column_index.index); + let filtered_col = filter(col, &bitmap)?; + + columns.push(filtered_col); + } + + Ok(Some(RecordBatch::try_new( + Arc::clone(output_schema), + columns, + )?)) + } + JoinType::RightMark | JoinType::LeftMark => { + if join_type == JoinType::RightMark { + debug_assert_eq!(batch_side, JoinSide::Right); + } + if join_type == JoinType::LeftMark { + debug_assert_eq!(batch_side, JoinSide::Left); + } + + let mut columns: Vec<Arc<dyn Array>> = + Vec::with_capacity(output_schema.fields().len()); + + // Hack to deal with the borrow checker + let mut right_batch_bitmap_opt = Some(batch_bitmap); + + for column_index in col_indices { + if column_index.side == batch_side { + let col = batch.column(column_index.index); + + columns.push(Arc::clone(col)); + } else if column_index.side == JoinSide::None { + let right_batch_bitmap = std::mem::take(&mut right_batch_bitmap_opt); + match right_batch_bitmap { + Some(right_batch_bitmap) => { + columns.push(Arc::new(right_batch_bitmap)) + } + None => unreachable!("Should only be one mark column"), + } + } else { + return internal_err!( + "Not possible to have this join side for RightMark join" + ); + } + } + + Ok(Some(RecordBatch::try_new( + Arc::clone(output_schema), + columns, + )?)) + } + _ => internal_err!( + "If batch is at right side, this function must be handling Full/Right/RightSemi/RightAnti/RightMark joins" + ), + } +} + +#[cfg(test)] +pub(crate) mod tests { + use super::*; + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::test::{TestMemoryExec, assert_join_metrics}; + use crate::{ + common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, + }; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::assert_contains; + use datafusion_common::test_util::batches_to_sort_string; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion_physical_expr::{Partitioning, PhysicalExpr}; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; + + use insta::allow_duplicates; + use insta::assert_snapshot; + use rstest::rstest; + + fn build_table( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), + batch_size: Option<usize>, + sorted_column_names: Vec<&str>, + ) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + + let batches = if let Some(batch_size) = batch_size { + let num_batches = batch.num_rows().div_ceil(batch_size); + (0..num_batches) + .map(|i| { + let start = i * batch_size; + let remaining_rows = batch.num_rows() - start; + batch.slice(start, batch_size.min(remaining_rows)) + }) + .collect::<Vec<_>>() + } else { + vec![batch] + }; + + let mut sort_info = vec![]; + for name in sorted_column_names { + let index = schema.index_of(name).unwrap(); + let sort_expr = PhysicalSortExpr::new( + Arc::new(Column::new(name, index)), + SortOptions::new(false, false), + ); + sort_info.push(sort_expr); + } + let mut source = TestMemoryExec::try_new(&[batches], schema, None).unwrap(); + if let Some(ordering) = LexOrdering::new(sort_info) { + source = source.try_with_sort_information(vec![ordering]).unwrap(); + } + + let source = Arc::new(source); + Arc::new(TestMemoryExec::update_cache(&source)) + } + + fn build_left_table() -> Arc<dyn ExecutionPlan> { + build_table( + ("a1", &vec![5, 9, 11]), + ("b1", &vec![5, 8, 8]), + ("c1", &vec![50, 90, 110]), + None, + Vec::new(), + ) + } + + fn build_right_table() -> Arc<dyn ExecutionPlan> { + build_table( + ("a2", &vec![12, 2, 10]), + ("b2", &vec![10, 2, 10]), + ("c2", &vec![40, 80, 100]), + None, + Vec::new(), + ) + } + + fn prepare_join_filter() -> JoinFilter { + let column_indices = vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let intermediate_schema = Schema::new(vec![ + Field::new("x", DataType::Int32, true), + Field::new("x", DataType::Int32, true), + ]); + // left.b1!=8 + let left_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(8)))), + )) as Arc<dyn PhysicalExpr>; + // right.b2!=10 + let right_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 1)), + Operator::NotEq, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + // filter = left.b1!=8 and right.b2!=10 + // after filter: + // left table: + // ("a1", &vec![5]), + // ("b1", &vec![5]), + // ("c1", &vec![50]), + // right table: + // ("a2", &vec![12, 2]), + // ("b2", &vec![10, 2]), + // ("c2", &vec![40, 80]), + let filter_expression = + Arc::new(BinaryExpr::new(left_filter, Operator::And, right_filter)) + as Arc<dyn PhysicalExpr>; + + JoinFilter::new( + filter_expression, + column_indices, + Arc::new(intermediate_schema), + ) + } + + pub(crate) async fn multi_partitioned_join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: &JoinType, + join_filter: Option<JoinFilter>, + context: Arc<TaskContext>, + ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> { + let partition_count = 4; + + // Redistributing right input + let right = Arc::new(RepartitionExec::try_new( + right, + Partitioning::RoundRobinBatch(partition_count), + )?) as Arc<dyn ExecutionPlan>; + + // Use the required distribution for nested loop join to test partition data + let nested_loop_join = + NestedLoopJoinExec::try_new(left, right, join_filter, join_type, None)?; + let columns = columns(&nested_loop_join.schema()); + let mut batches = vec![]; + for i in 0..partition_count { + let stream = nested_loop_join.execute(i, Arc::clone(&context))?; + let more_batches = common::collect(stream).await?; + batches.extend( + more_batches + .into_iter() + .inspect(|b| { + assert!(b.num_rows() <= context.session_config().batch_size()) + }) + .filter(|b| b.num_rows() > 0) + .collect::<Vec<_>>(), + ); + } + + let metrics = nested_loop_join.metrics().unwrap(); + + Ok((columns, batches, metrics)) + } + + fn new_task_ctx(batch_size: usize) -> Arc<TaskContext> { + let base = TaskContext::default(); + // limit max size of intermediate batch used in nlj to 1 + let cfg = base.session_config().clone().with_batch_size(batch_size); + Arc::new(base.with_session_config(cfg)) + } + + #[rstest] + #[tokio::test] + async fn join_inner_with_filter(#[values(1, 2, 16)] batch_size: usize) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + dbg!(&batch_size); + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::Inner, + Some(filter), + task_ctx, + ) + .await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 5 | 5 | 50 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + ")); + + assert_join_metrics!(metrics, 1); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_left_with_filter(#[values(1, 2, 16)] batch_size: usize) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::Left, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+-----+----+----+----+ + | 11 | 8 | 110 | | | | + | 5 | 5 | 50 | 2 | 2 | 80 | + | 9 | 8 | 90 | | | | + +----+----+-----+----+----+----+ + ")); + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_right_with_filter(#[values(1, 2, 16)] batch_size: usize) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::Right, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+-----+ + | | | | 10 | 10 | 100 | + | | | | 12 | 10 | 40 | + | 5 | 5 | 50 | 2 | 2 | 80 | + +----+----+----+----+----+-----+ + ")); + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_full_with_filter(#[values(1, 2, 16)] batch_size: usize) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::Full, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+-----+----+----+-----+ + | | | | 10 | 10 | 100 | + | | | | 12 | 10 | 40 | + | 11 | 8 | 110 | | | | + | 5 | 5 | 50 | 2 | 2 | 80 | + | 9 | 8 | 90 | | | | + +----+----+-----+----+----+-----+ + ")); + + assert_join_metrics!(metrics, 5); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_left_semi_with_filter( + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::LeftSemi, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 5 | 5 | 50 | + +----+----+----+ + ")); + + assert_join_metrics!(metrics, 1); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_left_anti_with_filter( + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::LeftAnti, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a1 | b1 | c1 | + +----+----+-----+ + | 11 | 8 | 110 | + | 9 | 8 | 90 | + +----+----+-----+ + ")); + + assert_join_metrics!(metrics, 2); + + Ok(()) + } + + #[tokio::test] + async fn join_has_correct_stats() -> Result<()> { + let left = build_left_table(); + let right = build_right_table(); + let nested_loop_join = NestedLoopJoinExec::try_new( + left, + right, + None, + &JoinType::Left, + Some(vec![1, 2]), + )?; + let stats = StatisticsContext::new() + .compute(&nested_loop_join, &StatisticsArgs::new())?; + assert_eq!( + nested_loop_join.schema().fields().len(), + stats.column_statistics.len(), + ); + assert_eq!(2, stats.column_statistics.len()); + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_right_semi_with_filter( + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::RightSemi, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a2", "b2", "c2"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+ + | a2 | b2 | c2 | + +----+----+----+ + | 2 | 2 | 80 | + +----+----+----+ + ")); + + assert_join_metrics!(metrics, 1); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_right_anti_with_filter( + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::RightAnti, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a2", "b2", "c2"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 10 | 10 | 100 | + | 12 | 10 | 40 | + +----+----+-----+ + ")); + + assert_join_metrics!(metrics, 2); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_left_mark_with_filter( + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::LeftMark, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a1", "b1", "c1", "mark"]); + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+-------+ + | a1 | b1 | c1 | mark | + +----+----+-----+-------+ + | 11 | 8 | 110 | false | + | 5 | 5 | 50 | true | + | 9 | 8 | 90 | false | + +----+----+-----+-------+ + ")); + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[rstest] + #[tokio::test] + async fn join_right_mark_with_filter( + #[values(1, 2, 16)] batch_size: usize, + ) -> Result<()> { + let task_ctx = new_task_ctx(batch_size); + let left = build_left_table(); + let right = build_right_table(); + + let filter = prepare_join_filter(); + let (columns, batches, metrics) = multi_partitioned_join_collect( + left, + right, + &JoinType::RightMark, + Some(filter), + task_ctx, + ) + .await?; + assert_eq!(columns, vec!["a2", "b2", "c2", "mark"]); + + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+-------+ + | a2 | b2 | c2 | mark | + +----+----+-----+-------+ + | 10 | 10 | 100 | false | + | 12 | 10 | 40 | false | + | 2 | 2 | 80 | true | + +----+----+-----+-------+ + ")); + + assert_join_metrics!(metrics, 3); + + Ok(()) + } + + #[tokio::test] + async fn test_overallocation() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("b1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + ("c1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), + None, + Vec::new(), + ); + let right = build_table( + ("a2", &vec![10, 11]), + ("b2", &vec![12, 13]), + ("c2", &vec![14, 15]), + None, + Vec::new(), + ); + let filter = prepare_join_filter(); + + // Join types that support memory-limited fallback should succeed + // even under tight memory limits (they spill to disk instead of OOM). + let fallback_join_types = vec![ + JoinType::Inner, + JoinType::Left, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::Right, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::RightMark, + ]; + + for join_type in &fallback_join_types { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // Should succeed via spill fallback, not OOM + let _result = multi_partitioned_join_collect( + Arc::clone(&left), + Arc::clone(&right), + join_type, + Some(filter.clone()), + task_ctx, + ) + .await?; + } + + // FULL JOIN with multiple right partitions is intentionally not + // supported in the fallback path yet (cross-partition left-bitmap + // coordination is missing). It should still OOM under tight memory. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + let err = multi_partitioned_join_collect( + Arc::clone(&left), + Arc::clone(&right), + &JoinType::Full, + Some(filter.clone()), + task_ctx, + ) + .await + .unwrap_err(); + assert_contains!(err.to_string(), "Resources exhausted"); + + Ok(()) + } + + /// Returns the column names on the schema + fn columns(schema: &Schema) -> Vec<String> { + schema.fields().iter().map(|f| f.name().clone()).collect() + } + + // ======================================================================== + // Memory-limited execution tests + // ======================================================================== + + /// Helper to run a NLJ using partition 0 and collect results + metrics. + async fn join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: &JoinType, + join_filter: Option<JoinFilter>, + context: Arc<TaskContext>, + ) -> Result<(Vec<String>, Vec<RecordBatch>, MetricsSet)> { + let nested_loop_join = + NestedLoopJoinExec::try_new(left, right, join_filter, join_type, None)?; + let columns = columns(&nested_loop_join.schema()); + let stream = nested_loop_join.execute(0, context)?; + let batches: Vec<RecordBatch> = common::collect(stream) + .await? + .into_iter() + .filter(|b| b.num_rows() > 0) + .collect(); + let metrics = nested_loop_join.metrics().unwrap(); + Ok((columns, batches, metrics)) + } + + /// Create a TaskContext with tight memory limit and disk spilling enabled. + fn task_ctx_with_memory_limit( + memory_limit: usize, + batch_size: usize, + ) -> Result<Arc<TaskContext>> { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .build_arc()?; + let cfg = TaskContext::default() + .session_config() + .clone() + .with_batch_size(batch_size); + let task_ctx = TaskContext::default() + .with_runtime(runtime) + .with_session_config(cfg); + Ok(Arc::new(task_ctx)) + } + + #[tokio::test] + async fn test_nlj_memory_limited_inner_join() -> Result<()> { + // Use a very small memory limit to force OOM → fallback to spill. + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::Inner, Some(filter), task_ctx).await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Verify spill actually occurred (memory-limited path was taken) + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + // Result should be identical to the non-memory-limited case + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 5 | 5 | 50 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_left_join() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::Left, Some(filter), task_ctx).await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Verify spill actually occurred + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+-----+----+----+----+ + | 11 | 8 | 110 | | | | + | 5 | 5 | 50 | 2 | 2 | 80 | + | 9 | 8 | 90 | | | | + +----+----+-----+----+----+----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_fits_in_memory_no_spill() -> Result<()> { + // Use a large memory limit — everything fits, no spilling needed. + let task_ctx = task_ctx_with_memory_limit(10_000_000, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::Inner, Some(filter), task_ctx).await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Verify no spilling occurred (standard OnceFut path was used) + assert_eq!( + metrics.spill_count().unwrap_or(0), + 0, + "Expected no spilling with generous memory limit" + ); + + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 5 | 5 | 50 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_empty_inputs() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + + // Empty left table + let empty_left = build_table( + ("a1", &vec![]), + ("b1", &vec![]), + ("c1", &vec![]), + None, + Vec::new(), + ); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (_columns, batches, _metrics) = + join_collect(empty_left, right, &JoinType::Inner, Some(filter), task_ctx) + .await?; + assert!(batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0)); + + // Empty right table + let task_ctx2 = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let empty_right = build_table( + ("a2", &vec![]), + ("b2", &vec![]), + ("c2", &vec![]), + None, + Vec::new(), + ); + let filter2 = prepare_join_filter(); + + let (_columns, batches, _metrics) = join_collect( + left, + empty_right, + &JoinType::Inner, + Some(filter2), + task_ctx2, + ) + .await?; + assert!(batches.is_empty() || batches.iter().all(|b| b.num_rows() == 0)); + + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_no_disk_falls_back_to_oom() -> Result<()> { + // When disk is disabled, fallback is not possible and OOM should occur. + use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); + + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let err = join_collect(left, right, &JoinType::Inner, Some(filter), task_ctx) + .await + .unwrap_err(); + + assert_contains!(err.to_string(), "Resources exhausted"); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_right_join() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::Right, Some(filter), task_ctx).await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Verify spill actually occurred + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + // Right join: all right rows appear. Unmatched right rows get NULLs on left. + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+-----+ + | | | | 10 | 10 | 100 | + | | | | 12 | 10 | 40 | + | 5 | 5 | 50 | 2 | 2 | 80 | + +----+----+----+----+----+-----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_full_join() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::Full, Some(filter), task_ctx).await?; + + assert_eq!(columns, vec!["a1", "b1", "c1", "a2", "b2", "c2"]); + + // Verify spill actually occurred + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + // Full join: unmatched from both sides appear with NULL padding. + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+----+----+-----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+-----+----+----+-----+ + | | | | 10 | 10 | 100 | + | | | | 12 | 10 | 40 | + | 11 | 8 | 110 | | | | + | 5 | 5 | 50 | 2 | 2 | 80 | + | 9 | 8 | 90 | | | | + +----+----+-----+----+----+-----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_right_semi_join() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::RightSemi, Some(filter), task_ctx) + .await?; + + assert_eq!(columns, vec!["a2", "b2", "c2"]); + + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + // Right semi: only right rows that matched at least one left row. + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+ + | a2 | b2 | c2 | + +----+----+----+ + | 2 | 2 | 80 | + +----+----+----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_right_anti_join() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::RightAnti, Some(filter), task_ctx) + .await?; + + assert_eq!(columns, vec!["a2", "b2", "c2"]); + + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + // Right anti: right rows that did NOT match any left row. + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+ + | a2 | b2 | c2 | + +----+----+-----+ + | 10 | 10 | 100 | + | 12 | 10 | 40 | + +----+----+-----+ + ")); + Ok(()) + } + + #[tokio::test] + async fn test_nlj_memory_limited_right_mark_join() -> Result<()> { + let task_ctx = task_ctx_with_memory_limit(50, 16)?; + let left = build_left_table(); + let right = build_right_table(); + let filter = prepare_join_filter(); + + let (columns, batches, metrics) = + join_collect(left, right, &JoinType::RightMark, Some(filter), task_ctx) + .await?; + + assert_eq!(columns, vec!["a2", "b2", "c2", "mark"]); + + assert!( + metrics.spill_count().unwrap_or(0) > 0, + "Expected spilling to occur under tight memory limit" + ); + + // Right mark: all right rows with a bool column indicating match. + allow_duplicates!(assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-----+-------+ + | a2 | b2 | c2 | mark | + +----+----+-----+-------+ + | 10 | 10 | 100 | false | + | 12 | 10 | 40 | false | + | 2 | 2 | 80 | true | + +----+----+-----+-------+ + ")); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/classic_join.rs new file mode 100644 index 0000000..50ef78f --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/classic_join.rs
@@ -0,0 +1,1546 @@ +// 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. + +//! Stream Implementation for PiecewiseMergeJoin's Classic Join (Left, Right, Full, Inner) + +use arrow::array::{Array, PrimitiveBuilder, new_null_array}; +use arrow::compute::{BatchCoalescer, take}; +use arrow::datatypes::UInt32Type; +use arrow::{ + array::{ArrayRef, RecordBatch, UInt32Array}, + compute::{sort_to_indices, take_record_batch}, +}; +use arrow_schema::{Schema, SchemaRef, SortOptions}; +use datafusion_common::NullEquality; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion_expr::{JoinType, Operator}; +use datafusion_physical_expr::PhysicalExprRef; +use futures::{Stream, StreamExt}; +use std::{cmp::Ordering, task::ready}; +use std::{sync::Arc, task::Poll}; + +use crate::handle_state; +use crate::joins::piecewise_merge_join::exec::{BufferedSide, BufferedSideReadyState}; +use crate::joins::piecewise_merge_join::utils::need_produce_result_in_final; +use crate::joins::utils::{BuildProbeJoinMetrics, StatefulStreamResult}; +use crate::joins::utils::{JoinKeyComparator, get_final_indices_from_shared_bitmap}; +use crate::stream::EmptyRecordBatchStream; + +pub(super) enum PiecewiseMergeJoinStreamState { + WaitBufferedSide, + FetchStreamBatch, + ProcessStreamBatch(SortedStreamBatch), + ProcessUnmatched, + Completed, +} + +impl PiecewiseMergeJoinStreamState { + // Grab mutable reference to the current stream batch + fn try_as_process_stream_batch_mut(&mut self) -> Result<&mut SortedStreamBatch> { + match self { + PiecewiseMergeJoinStreamState::ProcessStreamBatch(state) => Ok(state), + _ => internal_err!("Expected streamed batch in StreamBatch"), + } + } +} + +/// The stream side incoming batch with required sort order. +/// +/// Note the compare key in the join predicate might include expressions on the original +/// columns, so we store the evaluated compare key separately. +/// e.g. For join predicate `buffer.v1 < (stream.v1 + 1)`, the `compare_key_values` field stores +/// the evaluated `stream.v1 + 1` array. +pub(super) struct SortedStreamBatch { + pub batch: RecordBatch, + compare_key_values: Vec<ArrayRef>, +} + +impl SortedStreamBatch { + fn new(batch: RecordBatch, compare_key_values: Vec<ArrayRef>) -> Self { + Self { + batch, + compare_key_values, + } + } + + fn compare_key_values(&self) -> &Vec<ArrayRef> { + &self.compare_key_values + } +} + +pub(super) struct ClassicPWMJStream { + // Output schema of the `PiecewiseMergeJoin` + pub schema: Arc<Schema>, + + // Physical expression that is evaluated on the streamed side + // We do not need on_buffered as this is already evaluated when + // creating the buffered side which happens before initializing + // `PiecewiseMergeJoinStream` + pub on_streamed: PhysicalExprRef, + // Type of join + pub join_type: JoinType, + // Comparison operator + pub operator: Operator, + // Streamed batch + pub streamed: SendableRecordBatchStream, + // Streamed schema + streamed_schema: SchemaRef, + // Buffered side data + buffered_side: BufferedSide, + // Tracks the state of the `PiecewiseMergeJoin` + state: PiecewiseMergeJoinStreamState, + // Sort option for streamed side (specifies whether + // the sort is ascending or descending) + sort_option: SortOptions, + // Metrics for build + probe joins + join_metrics: BuildProbeJoinMetrics, + // Tracking incremental state for emitting record batches + batch_process_state: BatchProcessState, +} + +impl RecordBatchStream for ClassicPWMJStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +// `PiecewiseMergeJoinStreamState` is separated into `WaitBufferedSide`, `FetchStreamBatch`, +// `ProcessStreamBatch`, `ProcessUnmatched` and `Completed`. +// +// Classic Joins +// 1. `WaitBufferedSide` - Load in the buffered side data into memory. +// 2. `FetchStreamBatch` - Fetch + sort incoming stream batches. We switch the state to +// `Completed` if there are still remaining partitions to process. It is only switched to +// `ExhaustedStreamBatch` if all partitions have been processed. +// 3. `ProcessStreamBatch` - Compare stream batch row values against the buffered side data. +// 4. `ExhaustedStreamBatch` - If the join type is Left or Inner we will return state as +// `Completed` however for Full and Right we will need to process the unmatched buffered rows. +impl ClassicPWMJStream { + // Creates a new `PiecewiseMergeJoinStream` instance + #[expect(clippy::too_many_arguments)] + pub fn try_new( + schema: Arc<Schema>, + on_streamed: PhysicalExprRef, + join_type: JoinType, + operator: Operator, + streamed: SendableRecordBatchStream, + buffered_side: BufferedSide, + state: PiecewiseMergeJoinStreamState, + sort_option: SortOptions, + join_metrics: BuildProbeJoinMetrics, + batch_size: usize, + ) -> Self { + Self { + schema: Arc::clone(&schema), + on_streamed, + join_type, + operator, + streamed_schema: streamed.schema(), + streamed, + buffered_side, + state, + sort_option, + join_metrics, + batch_process_state: BatchProcessState::new(schema, batch_size), + } + } + + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + return match self.state { + PiecewiseMergeJoinStreamState::WaitBufferedSide => { + handle_state!(ready!(self.collect_buffered_side(cx))) + } + PiecewiseMergeJoinStreamState::FetchStreamBatch => { + handle_state!(ready!(self.fetch_stream_batch(cx))) + } + PiecewiseMergeJoinStreamState::ProcessStreamBatch(_) => { + handle_state!(self.process_stream_batch()) + } + PiecewiseMergeJoinStreamState::ProcessUnmatched => { + handle_state!(self.process_unmatched_buffered_batch()) + } + PiecewiseMergeJoinStreamState::Completed => Poll::Ready(None), + }; + } + } + + // Collects buffered side data + fn collect_buffered_side( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + let build_timer = self.join_metrics.build_time.timer(); + let buffered_data = ready!( + self.buffered_side + .try_as_initial_mut()? + .buffered_fut + .get_shared(cx) + )?; + build_timer.done(); + + // We will start fetching stream batches for classic joins + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + + self.buffered_side = + BufferedSide::Ready(BufferedSideReadyState { buffered_data }); + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + // Fetches incoming stream batches + fn fetch_stream_batch( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.streamed.poll_next_unpin(cx)) { + None => { + // Release the streamed input pipeline's resources. + let streamed_schema = self.streamed.schema(); + self.streamed = Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + if self + .buffered_side + .try_as_ready_mut()? + .buffered_data + .remaining_partitions + .fetch_sub(1, std::sync::atomic::Ordering::SeqCst) + == 1 + { + self.batch_process_state.reset(); + self.state = PiecewiseMergeJoinStreamState::ProcessUnmatched; + } else { + self.state = PiecewiseMergeJoinStreamState::Completed; + } + } + Some(Ok(batch)) => { + // Evaluate the streamed physical expression on the stream batch + let stream_values: ArrayRef = self + .on_streamed + .evaluate(&batch)? + .into_array(batch.num_rows())?; + + self.join_metrics.input_batches.add(1); + self.join_metrics.input_rows.add(batch.num_rows()); + + // Sort stream values and change the streamed record batch accordingly + let indices = sort_to_indices( + stream_values.as_ref(), + Some(self.sort_option), + None, + )?; + let stream_batch = take_record_batch(&batch, &indices)?; + let stream_values = take(stream_values.as_ref(), &indices, None)?; + + // Reset BatchProcessState before processing a new stream batch + self.batch_process_state.reset(); + self.state = PiecewiseMergeJoinStreamState::ProcessStreamBatch( + SortedStreamBatch::new(stream_batch, vec![stream_values]), + ); + } + Some(Err(err)) => return Poll::Ready(Err(err)), + }; + + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + + // Only classic join will call. This function will process stream batches and evaluate against + // the buffered side data. + fn process_stream_batch( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + let buffered_side = self.buffered_side.try_as_ready_mut()?; + let stream_batch = self.state.try_as_process_stream_batch_mut()?; + + if let Some(batch) = self + .batch_process_state + .output_batches + .next_completed_batch() + { + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + + // Produce more work + let batch = resolve_classic_join( + buffered_side, + stream_batch, + &self.schema, + self.operator, + self.sort_option, + self.join_type, + &mut self.batch_process_state, + )?; + + if !self.batch_process_state.continue_process { + // We finished scanning this stream batch. + self.batch_process_state + .output_batches + .finish_buffered_batch()?; + if let Some(b) = self + .batch_process_state + .output_batches + .next_completed_batch() + { + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + return Ok(StatefulStreamResult::Ready(Some(b))); + } + + // Nothing pending; hand back whatever `resolve` returned (often empty) and move on. + if self.batch_process_state.output_batches.is_empty() { + self.state = PiecewiseMergeJoinStreamState::FetchStreamBatch; + + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + } + + Ok(StatefulStreamResult::Ready(Some(batch))) + } + + // Process remaining unmatched rows + fn process_unmatched_buffered_batch( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + // Return early for `JoinType::Right` and `JoinType::Inner` + if matches!(self.join_type, JoinType::Right | JoinType::Inner) { + self.state = PiecewiseMergeJoinStreamState::Completed; + return Ok(StatefulStreamResult::Ready(None)); + } + + if !self.batch_process_state.continue_process { + if let Some(batch) = self + .batch_process_state + .output_batches + .next_completed_batch() + { + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + + self.batch_process_state + .output_batches + .finish_buffered_batch()?; + if let Some(batch) = self + .batch_process_state + .output_batches + .next_completed_batch() + { + self.state = PiecewiseMergeJoinStreamState::Completed; + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + } + + let buffered_data = + Arc::clone(&self.buffered_side.try_as_ready().unwrap().buffered_data); + + let (buffered_indices, _streamed_indices) = get_final_indices_from_shared_bitmap( + &buffered_data.visited_indices_bitmap, + self.join_type, + true, + ); + + let new_buffered_batch = + take_record_batch(buffered_data.batch(), &buffered_indices)?; + let mut buffered_columns = new_buffered_batch.columns().to_vec(); + + let streamed_columns: Vec<ArrayRef> = self + .streamed_schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), new_buffered_batch.num_rows())) + .collect(); + + buffered_columns.extend(streamed_columns); + + let batch = RecordBatch::try_new(Arc::clone(&self.schema), buffered_columns)?; + + self.batch_process_state.output_batches.push_batch(batch)?; + + self.batch_process_state.continue_process = false; + if let Some(batch) = self + .batch_process_state + .output_batches + .next_completed_batch() + { + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + + self.batch_process_state + .output_batches + .finish_buffered_batch()?; + if let Some(batch) = self + .batch_process_state + .output_batches + .next_completed_batch() + { + self.state = PiecewiseMergeJoinStreamState::Completed; + return Ok(StatefulStreamResult::Ready(Some(batch))); + } + + self.state = PiecewiseMergeJoinStreamState::Completed; + self.batch_process_state.reset(); + Ok(StatefulStreamResult::Ready(None)) + } +} + +struct BatchProcessState { + // Used to pick up from the last index on the stream side + output_batches: Box<BatchCoalescer>, + // Used to store the unmatched stream indices for `JoinType::Right` and `JoinType::Full` + unmatched_indices: PrimitiveBuilder<UInt32Type>, + // Used to store the start index on the buffered side; used to resume processing on the correct + // row + start_buffer_idx: usize, + // Used to store the start index on the stream side; used to resume processing on the correct + // row + start_stream_idx: usize, + // Signals if we found a match for the current stream row + found: bool, + // Signals to continue processing the current stream batch + continue_process: bool, + // Skip nulls + processed_null_count: bool, +} + +impl BatchProcessState { + pub(crate) fn new(schema: Arc<Schema>, batch_size: usize) -> Self { + Self { + output_batches: Box::new(BatchCoalescer::new(schema, batch_size)), + unmatched_indices: PrimitiveBuilder::new(), + start_buffer_idx: 0, + start_stream_idx: 0, + found: false, + continue_process: true, + processed_null_count: false, + } + } + + pub(crate) fn reset(&mut self) { + self.unmatched_indices = PrimitiveBuilder::new(); + self.start_buffer_idx = 0; + self.start_stream_idx = 0; + self.found = false; + self.continue_process = true; + self.processed_null_count = false; + } +} + +impl Stream for ClassicPWMJStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_impl(cx) + } +} + +// For Left, Right, Full, and Inner joins, incoming stream batches will already be sorted. +fn resolve_classic_join( + buffered_side: &mut BufferedSideReadyState, + stream_batch: &SortedStreamBatch, + join_schema: &SchemaRef, + operator: Operator, + sort_options: SortOptions, + join_type: JoinType, + batch_process_state: &mut BatchProcessState, +) -> Result<RecordBatch> { + let buffered_len = buffered_side.buffered_data.values().len(); + let stream_values = stream_batch.compare_key_values(); + + // Build comparator once for the batch pair + let cmp = JoinKeyComparator::new( + &[Arc::clone(&stream_values[0])], + &[Arc::clone(buffered_side.buffered_data.values())], + &[sort_options], + NullEquality::NullEqualsNothing, + )?; + + let mut buffer_idx = batch_process_state.start_buffer_idx; + let mut stream_idx = batch_process_state.start_stream_idx; + + if !batch_process_state.processed_null_count { + let buffered_null_idx = buffered_side.buffered_data.values().null_count(); + let stream_null_idx = stream_values[0].null_count(); + buffer_idx = buffered_null_idx; + stream_idx = stream_null_idx; + batch_process_state.processed_null_count = true; + } + + // Our buffer_idx variable allows us to start probing on the buffered side where we last matched + // in the previous stream row. + for row_idx in stream_idx..stream_batch.batch.num_rows() { + while buffer_idx < buffered_len { + let compare = cmp.compare(row_idx, buffer_idx); + + // If we find a match we append all indices and move to the next stream row index + match operator { + Operator::Gt | Operator::Lt => { + if compare == Ordering::Less { + batch_process_state.found = true; + let count = buffered_len - buffer_idx; + + let batch = build_matched_indices_and_set_buffered_bitmap( + (buffer_idx, count), + (row_idx, count), + buffered_side, + stream_batch, + join_type, + join_schema, + )?; + + batch_process_state.output_batches.push_batch(batch)?; + + // Flush batch and update pointers if we have a completed batch + if let Some(batch) = + batch_process_state.output_batches.next_completed_batch() + { + batch_process_state.found = false; + batch_process_state.start_buffer_idx = buffer_idx; + batch_process_state.start_stream_idx = row_idx + 1; + return Ok(batch); + } + + break; + } + } + Operator::GtEq | Operator::LtEq => { + if matches!(compare, Ordering::Equal | Ordering::Less) { + batch_process_state.found = true; + let count = buffered_len - buffer_idx; + let batch = build_matched_indices_and_set_buffered_bitmap( + (buffer_idx, count), + (row_idx, count), + buffered_side, + stream_batch, + join_type, + join_schema, + )?; + + // Flush batch and update pointers if we have a completed batch + batch_process_state.output_batches.push_batch(batch)?; + if let Some(batch) = + batch_process_state.output_batches.next_completed_batch() + { + batch_process_state.found = false; + batch_process_state.start_buffer_idx = buffer_idx; + batch_process_state.start_stream_idx = row_idx + 1; + return Ok(batch); + } + + break; + } + } + _ => { + return internal_err!( + "PiecewiseMergeJoin should not contain operator, {}", + operator + ); + } + }; + + // Increment buffer_idx after every row + buffer_idx += 1; + } + + // If a match was not found for the current stream row index the stream indice is appended + // to the unmatched indices to be flushed later. + if matches!(join_type, JoinType::Right | JoinType::Full) + && !batch_process_state.found + { + batch_process_state + .unmatched_indices + .append_value(row_idx as u32); + } + + batch_process_state.found = false; + } + + // Flushed all unmatched indices on the streamed side + if matches!(join_type, JoinType::Right | JoinType::Full) { + let batch = create_unmatched_batch( + &mut batch_process_state.unmatched_indices, + stream_batch, + join_schema, + )?; + + batch_process_state.output_batches.push_batch(batch)?; + } + + batch_process_state.continue_process = false; + Ok(RecordBatch::new_empty(Arc::clone(join_schema))) +} + +// Builds a record batch from indices ranges on the buffered and streamed side. +// +// The two ranges are: buffered_range: (start index, count) and streamed_range: (start index, count) due +// to batch.slice(start, count). +fn build_matched_indices_and_set_buffered_bitmap( + buffered_range: (usize, usize), + streamed_range: (usize, usize), + buffered_side: &mut BufferedSideReadyState, + stream_batch: &SortedStreamBatch, + join_type: JoinType, + join_schema: &SchemaRef, +) -> Result<RecordBatch> { + // Mark the buffered indices as visited + if need_produce_result_in_final(join_type) { + let mut bitmap = buffered_side.buffered_data.visited_indices_bitmap.lock(); + for i in buffered_range.0..buffered_range.0 + buffered_range.1 { + bitmap.set_bit(i, true); + } + } + + let new_buffered_batch = buffered_side + .buffered_data + .batch() + .slice(buffered_range.0, buffered_range.1); + let mut buffered_columns = new_buffered_batch.columns().to_vec(); + + let indices = UInt32Array::from_value(streamed_range.0 as u32, streamed_range.1); + let new_stream_batch = take_record_batch(&stream_batch.batch, &indices)?; + let streamed_columns = new_stream_batch.columns().to_vec(); + + buffered_columns.extend(streamed_columns); + + Ok(RecordBatch::try_new( + Arc::clone(join_schema), + buffered_columns, + )?) +} + +// Creates a record batch from the unmatched indices on the streamed side +fn create_unmatched_batch( + streamed_indices: &mut PrimitiveBuilder<UInt32Type>, + stream_batch: &SortedStreamBatch, + join_schema: &SchemaRef, +) -> Result<RecordBatch> { + let streamed_indices = streamed_indices.finish(); + let new_stream_batch = take_record_batch(&stream_batch.batch, &streamed_indices)?; + let streamed_columns = new_stream_batch.columns().to_vec(); + let buffered_cols_len = join_schema.fields().len() - streamed_columns.len(); + + let num_rows = new_stream_batch.num_rows(); + let mut buffered_columns: Vec<ArrayRef> = join_schema + .fields() + .iter() + .take(buffered_cols_len) + .map(|field| new_null_array(field.data_type(), num_rows)) + .collect(); + + buffered_columns.extend(streamed_columns); + + Ok(RecordBatch::try_new( + Arc::clone(join_schema), + buffered_columns, + )?) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + ExecutionPlan, common, + joins::PiecewiseMergeJoinExec, + test::{TestMemoryExec, build_table_i32}, + }; + use arrow::array::{Date32Array, Date64Array}; + use arrow_schema::{DataType, Field}; + use datafusion_common::test_util::batches_to_string; + use datafusion_execution::TaskContext; + use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; + use insta::assert_snapshot; + use std::sync::Arc; + + fn columns(schema: &Schema) -> Vec<String> { + schema.fields().iter().map(|f| f.name().clone()).collect() + } + + fn build_table( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), + ) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + fn build_date_table( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), + ) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Date32, false), + Field::new(b.0, DataType::Date32, false), + Field::new(c.0, DataType::Date32, false), + ]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Date32Array::from(a.1.clone())), + Arc::new(Date32Array::from(b.1.clone())), + Arc::new(Date32Array::from(c.1.clone())), + ], + ) + .unwrap(); + + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + fn build_date64_table( + a: (&str, &Vec<i64>), + b: (&str, &Vec<i64>), + c: (&str, &Vec<i64>), + ) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Date64, false), + Field::new(b.0, DataType::Date64, false), + Field::new(c.0, DataType::Date64, false), + ]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Date64Array::from(a.1.clone())), + Arc::new(Date64Array::from(b.1.clone())), + Arc::new(Date64Array::from(c.1.clone())), + ], + ) + .unwrap(); + + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + fn join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>), + operator: Operator, + join_type: JoinType, + ) -> Result<PiecewiseMergeJoinExec> { + PiecewiseMergeJoinExec::try_new(left, right, on, operator, join_type, 1) + } + + async fn join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: (PhysicalExprRef, PhysicalExprRef), + operator: Operator, + join_type: JoinType, + ) -> Result<(Vec<String>, Vec<RecordBatch>)> { + join_collect_with_options(left, right, on, operator, join_type).await + } + + async fn join_collect_with_options( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: (PhysicalExprRef, PhysicalExprRef), + operator: Operator, + join_type: JoinType, + ) -> Result<(Vec<String>, Vec<RecordBatch>)> { + let task_ctx = Arc::new(TaskContext::default()); + let join = join(left, right, on, operator, join_type)?; + let columns = columns(&join.schema()); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + Ok((columns, batches)) + } + + #[tokio::test] + async fn join_inner_less_than() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 3 | 7 | + // | 2 | 2 | 8 | + // | 3 | 1 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![3, 2, 1]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 2 | 70 | + // | 20 | 3 | 80 | + // | 30 | 4 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 3 | 7 | 30 | 4 | 90 | + | 2 | 2 | 8 | 30 | 4 | 90 | + | 3 | 1 | 9 | 30 | 4 | 90 | + | 2 | 2 | 8 | 20 | 3 | 80 | + | 3 | 1 | 9 | 20 | 3 | 80 | + | 3 | 1 | 9 | 10 | 2 | 70 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_less_than_unsorted() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 3 | 7 | + // | 2 | 2 | 8 | + // | 3 | 1 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![3, 2, 1]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 3 | 70 | + // | 20 | 2 | 80 | + // | 30 | 4 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![3, 2, 4]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 3 | 7 | 30 | 4 | 90 | + | 2 | 2 | 8 | 30 | 4 | 90 | + | 3 | 1 | 9 | 30 | 4 | 90 | + | 2 | 2 | 8 | 10 | 3 | 70 | + | 3 | 1 | 9 | 10 | 3 | 70 | + | 3 | 1 | 9 | 20 | 2 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_greater_than_equal_to() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 2 | 7 | + // | 2 | 3 | 8 | + // | 3 | 4 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![2, 3, 4]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 3 | 70 | + // | 20 | 2 | 80 | + // | 30 | 1 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![3, 2, 1]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::GtEq, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 2 | 7 | 30 | 1 | 90 | + | 2 | 3 | 8 | 30 | 1 | 90 | + | 3 | 4 | 9 | 30 | 1 | 90 | + | 1 | 2 | 7 | 20 | 2 | 80 | + | 2 | 3 | 8 | 20 | 2 | 80 | + | 3 | 4 | 9 | 20 | 2 | 80 | + | 2 | 3 | 8 | 10 | 3 | 70 | + | 3 | 4 | 9 | 10 | 3 | 70 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_empty_left() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // (empty) + // +----+----+----+ + let left = build_table( + ("a1", &Vec::<i32>::new()), + ("b1", &Vec::<i32>::new()), + ("c1", &Vec::<i32>::new()), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 1 | 1 | 1 | + // | 2 | 2 | 2 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![1, 2]), + ("b1", &vec![1, 2]), + ("c2", &vec![1, 2]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + let (_, batches) = + join_collect(left, right, on, Operator::LtEq, JoinType::Inner).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_full_greater_than_equal_to() -> Result<()> { + // +----+----+-----+ + // | a1 | b1 | c1 | + // +----+----+-----+ + // | 1 | 1 | 100 | + // | 2 | 2 | 200 | + // +----+----+-----+ + let left = build_table( + ("a1", &vec![1, 2]), + ("b1", &vec![1, 2]), + ("c1", &vec![100, 200]), + ); + + // +----+----+-----+ + // | a2 | b1 | c2 | + // +----+----+-----+ + // | 10 | 3 | 300 | + // | 20 | 2 | 400 | + // +----+----+-----+ + let right = build_table( + ("a2", &vec![10, 20]), + ("b1", &vec![3, 2]), + ("c2", &vec![300, 400]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::GtEq, JoinType::Full).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+----+----+-----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+-----+----+----+-----+ + | 2 | 2 | 200 | 20 | 2 | 400 | + | | | | 10 | 3 | 300 | + | 1 | 1 | 100 | | | | + +----+----+-----+----+----+-----+ + "); + + Ok(()) + } + + #[tokio::test] + async fn join_left_greater_than() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 1 | 7 | + // | 2 | 3 | 8 | + // | 3 | 4 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 3, 4]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 3 | 70 | + // | 20 | 2 | 80 | + // | 30 | 1 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![3, 2, 1]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::Left).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 2 | 3 | 8 | 30 | 1 | 90 | + | 3 | 4 | 9 | 30 | 1 | 90 | + | 2 | 3 | 8 | 20 | 2 | 80 | + | 3 | 4 | 9 | 20 | 2 | 80 | + | 3 | 4 | 9 | 10 | 3 | 70 | + | 1 | 1 | 7 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_right_greater_than() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 1 | 7 | + // | 2 | 3 | 8 | + // | 3 | 4 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 3, 4]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 5 | 70 | + // | 20 | 3 | 80 | + // | 30 | 2 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![5, 3, 2]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::Right).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 2 | 3 | 8 | 30 | 2 | 90 | + | 3 | 4 | 9 | 30 | 2 | 90 | + | 3 | 4 | 9 | 20 | 3 | 80 | + | | | | 10 | 5 | 70 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_right_less_than() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 4 | 7 | + // | 2 | 3 | 8 | + // | 3 | 1 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 3, 1]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 2 | 70 | + // | 20 | 3 | 80 | + // | 30 | 5 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![2, 3, 5]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Right).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 30 | 5 | 90 | + | 2 | 3 | 8 | 30 | 5 | 90 | + | 3 | 1 | 9 | 30 | 5 | 90 | + | 3 | 1 | 9 | 20 | 3 | 80 | + | 3 | 1 | 9 | 10 | 2 | 70 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_less_than_equal_with_dups() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 4 | 7 | + // | 2 | 4 | 8 | + // | 3 | 2 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 4, 2]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 4 | 70 | + // | 20 | 3 | 80 | + // | 30 | 2 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 3, 2]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::LtEq, JoinType::Inner).await?; + + // Expected grouping follows right.b1 descending (4, 3, 2) + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 4 | 8 | 10 | 4 | 70 | + | 3 | 2 | 9 | 10 | 4 | 70 | + | 3 | 2 | 9 | 20 | 3 | 80 | + | 3 | 2 | 9 | 30 | 2 | 90 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_greater_than_unsorted_right() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 1 | 7 | + // | 2 | 2 | 8 | + // | 3 | 4 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 4]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 3 | 70 | + // | 20 | 1 | 80 | + // | 30 | 2 | 90 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![3, 1, 2]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::Inner).await?; + + // Grouped by right in ascending evaluation for > (1,2,3) + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 2 | 2 | 8 | 20 | 1 | 80 | + | 3 | 4 | 9 | 20 | 1 | 80 | + | 3 | 4 | 9 | 30 | 2 | 90 | + | 3 | 4 | 9 | 10 | 3 | 70 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_left_less_than_equal_with_left_nulls_on_no_match() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 5 | 7 | + // | 2 | 4 | 8 | + // | 3 | 1 | 9 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![5, 4, 1]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 3 | 70 | + // +----+----+----+ + let right = build_table(("a2", &vec![10]), ("b1", &vec![3]), ("c2", &vec![70])); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::LtEq, JoinType::Left).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 3 | 1 | 9 | 10 | 3 | 70 | + | 1 | 5 | 7 | | | | + | 2 | 4 | 8 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_right_greater_than_equal_with_right_nulls_on_no_match() -> Result<()> { + // +----+----+----+ + // | a1 | b1 | c1 | + // +----+----+----+ + // | 1 | 1 | 7 | + // | 2 | 2 | 8 | + // +----+----+----+ + let left = build_table( + ("a1", &vec![1, 2]), + ("b1", &vec![1, 2]), + ("c1", &vec![7, 8]), + ); + + // +----+----+----+ + // | a2 | b1 | c2 | + // +----+----+----+ + // | 10 | 3 | 70 | + // | 20 | 5 | 80 | + // +----+----+----+ + let right = build_table( + ("a2", &vec![10, 20]), + ("b1", &vec![3, 5]), + ("c2", &vec![70, 80]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::GtEq, JoinType::Right).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | | | | 10 | 3 | 70 | + | | | | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_single_row_left_less_than() -> Result<()> { + let left = build_table(("a1", &vec![42]), ("b1", &vec![5]), ("c1", &vec![999])); + + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![1, 5, 7]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+-----+----+----+----+ + | 42 | 5 | 999 | 30 | 7 | 90 | + +----+----+-----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_inner_empty_right() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1, 2, 3]), + ("c1", &vec![7, 8, 9]), + ); + + let right = build_table( + ("a2", &Vec::<i32>::new()), + ("b1", &Vec::<i32>::new()), + ("c2", &Vec::<i32>::new()), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Gt, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + +----+----+----+----+----+----+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_date32_inner_less_than() -> Result<()> { + // +----+-------+----+ + // | a1 | b1 | c1 | + // +----+-------+----+ + // | 1 | 19107 | 7 | + // | 2 | 19107 | 8 | + // | 3 | 19105 | 9 | + // +----+-------+----+ + let left = build_date_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![19107, 19107, 19105]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+-------+----+ + // | a2 | b1 | c2 | + // +----+-------+----+ + // | 10 | 19105 | 70 | + // | 20 | 19103 | 80 | + // | 30 | 19107 | 90 | + // +----+-------+----+ + let right = build_date_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![19105, 19103, 19107]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +------------+------------+------------+------------+------------+------------+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +------------+------------+------------+------------+------------+------------+ + | 1970-01-04 | 2022-04-23 | 1970-01-10 | 1970-01-31 | 2022-04-25 | 1970-04-01 | + +------------+------------+------------+------------+------------+------------+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_date64_inner_less_than() -> Result<()> { + // +----+---------------+----+ + // | a1 | b1 | c1 | + // +----+---------------+----+ + // | 1 | 1650903441000 | 7 | + // | 2 | 1650903441000 | 8 | + // | 3 | 1650703441000 | 9 | + // +----+---------------+----+ + let left = build_date64_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1650903441000, 1650903441000, 1650703441000]), + ("c1", &vec![7, 8, 9]), + ); + + // +----+---------------+----+ + // | a2 | b1 | c2 | + // +----+---------------+----+ + // | 10 | 1650703441000 | 70 | + // | 20 | 1650503441000 | 80 | + // | 30 | 1650903441000 | 90 | + // +----+---------------+----+ + let right = build_date64_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![1650703441000, 1650503441000, 1650903441000]), + ("c2", &vec![70, 80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Inner).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + | 1970-01-01T00:00:00.003 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.009 | 1970-01-01T00:00:00.030 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.090 | + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + "); + Ok(()) + } + + #[tokio::test] + async fn join_date64_right_less_than() -> Result<()> { + // +----+---------------+----+ + // | a1 | b1 | c1 | + // +----+---------------+----+ + // | 1 | 1650903441000 | 7 | + // | 2 | 1650703441000 | 8 | + // +----+---------------+----+ + let left = build_date64_table( + ("a1", &vec![1, 2]), + ("b1", &vec![1650903441000, 1650703441000]), + ("c1", &vec![7, 8]), + ); + + // +----+---------------+----+ + // | a2 | b1 | c2 | + // +----+---------------+----+ + // | 10 | 1650703441000 | 80 | + // | 20 | 1650903441000 | 90 | + // +----+---------------+----+ + let right = build_date64_table( + ("a2", &vec![10, 20]), + ("b1", &vec![1650703441000, 1650903441000]), + ("c2", &vec![80, 90]), + ); + + let on = ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ); + + let (_, batches) = + join_collect(left, right, on, Operator::Lt, JoinType::Right).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + | 1970-01-01T00:00:00.002 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.008 | 1970-01-01T00:00:00.020 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.090 | + | | | | 1970-01-01T00:00:00.010 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.080 | + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + "); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/exec.rs b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/exec.rs new file mode 100644 index 0000000..c42ec67 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/exec.rs
@@ -0,0 +1,819 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::Array; +use arrow::{ + array::{ArrayRef, BooleanBufferBuilder, RecordBatch}, + compute::concat_batches, + util::bit_util, +}; +use arrow_schema::{SchemaRef, SortOptions}; +use datafusion_common::not_impl_err; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{JoinSide, Result, internal_err}; +use datafusion_execution::{ + SendableRecordBatchStream, + memory_pool::{MemoryConsumer, MemoryReservation}, +}; +use datafusion_expr::{JoinType, Operator}; +use datafusion_physical_expr::equivalence::join_equivalence_properties; +use datafusion_physical_expr::{ + Distribution, LexOrdering, OrderingRequirements, PhysicalExpr, PhysicalExprRef, + PhysicalSortExpr, +}; +use datafusion_physical_expr_common::physical_expr::fmt_sql; +use futures::TryStreamExt; +use parking_lot::Mutex; +use std::fmt::Formatter; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; + +use crate::execution_plan::{EmissionType, boundedness_from_children}; + +use crate::joins::piecewise_merge_join::classic_join::{ + ClassicPWMJStream, PiecewiseMergeJoinStreamState, +}; +use crate::joins::piecewise_merge_join::utils::{ + build_visited_indices_map, is_existence_join, is_right_existence_join, +}; +use crate::joins::utils::asymmetric_join_output_partitioning; +use crate::metrics::MetricsSet; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + ReplaceChildrenOptions, validate_child_count, +}; +use crate::{ + ExecutionPlan, PlanProperties, + joins::{ + SharedBitmapBuilder, + utils::{BuildProbeJoinMetrics, OnceAsync, OnceFut, build_join_schema}, + }, + metrics::ExecutionPlanMetricsSet, + spill::get_record_batch_memory_size, +}; + +/// `PiecewiseMergeJoinExec` is a join execution plan that only evaluates single range filter and show much +/// better performance for these workloads than `NestedLoopJoin` +/// +/// The physical planner will choose to evaluate this join when there is only one comparison filter. This +/// is a binary expression which contains [`Operator::Lt`], [`Operator::LtEq`], [`Operator::Gt`], and +/// [`Operator::GtEq`].: +/// Examples: +/// - `col0` < `colb`, `col0` <= `colb`, `col0` > `colb`, `col0` >= `colb` +/// +/// # Execution Plan Inputs +/// For `PiecewiseMergeJoin` we label all right inputs as the `streamed' side and the left outputs as the +/// 'buffered' side. +/// +/// `PiecewiseMergeJoin` takes a sorted input for the side to be buffered and is able to sort streamed record +/// batches during processing. Sorted input must specifically be ascending/descending based on the operator. +/// +/// # Algorithms +/// Classic joins are processed differently compared to existence joins. +/// +/// ## Classic Joins (Inner, Full, Left, Right) +/// For classic joins we buffer the build side and stream the probe side (the "probe" side). +/// Both sides are sorted so that we can iterate from index 0 to the end on each side. This ordering ensures +/// that when we find the first matching pair of rows, we can emit the current stream row joined with all remaining +/// probe rows from the match position onward, without rescanning earlier probe rows. +/// +/// For `<` and `<=` operators, both inputs are sorted in **descending** order, while for `>` and `>=` operators +/// they are sorted in **ascending** order. This choice ensures that the pointer on the buffered side can advance +/// monotonically as we stream new batches from the stream side. +/// +/// The streamed side may arrive unsorted, so this operator sorts each incoming batch in memory before +/// processing. The buffered side is required to be globally sorted; the plan declares this requirement +/// in `requires_input_order`, which allows the optimizer to automatically insert a `SortExec` on that side if needed. +/// By the time this operator runs, the buffered side is guaranteed to be in the proper order. +/// +/// The pseudocode for the algorithm looks like this: +/// +/// ```text +/// for stream_row in stream_batch: +/// for buffer_row in buffer_batch: +/// if compare(stream_row, probe_row): +/// output stream_row X buffer_batch[buffer_row:] +/// else: +/// continue +/// ``` +/// +/// The algorithm uses the streamed side (larger) to drive the loop. This is due to every row on the stream side iterating +/// the buffered side to find every first match. By doing this, each match can output more result so that output +/// handling can be better vectorized for performance. +/// +/// Here is an example: +/// +/// We perform a `JoinType::Left` with these two batches and the operator being `Operator::Lt`(<). For each +/// row on the streamed side we move a pointer on the buffered until it matches the condition. Once we reach +/// the row which matches (in this case with row 1 on streamed will have its first match on row 2 on +/// buffered; 100 < 200 is true), we can emit all rows after that match. We can emit the rows like this because +/// if the batch is sorted in ascending order, every subsequent row will also satisfy the condition as they will +/// all be larger values. +/// +/// ```text +/// SQL statement: +/// SELECT * +/// FROM (VALUES (100), (200), (500)) AS streamed(a) +/// LEFT JOIN (VALUES (100), (200), (200), (300), (400)) AS buffered(b) +/// ON streamed.a < buffered.b; +/// +/// Processing Row 1: +/// +/// Sorted Buffered Side Sorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 100 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ ─┐ 2 │ 200 │ +/// ├──────────────────┤ │ For row 1 on streamed side with ├──────────────────┤ +/// 3 │ 200 │ │ value 100, we emit rows 2 - 5. 3 │ 500 │ +/// ├──────────────────┤ │ as matches when the operator is └──────────────────┘ +/// 4 │ 300 │ │ `Operator::Lt` (<) Emitting all +/// ├──────────────────┤ │ rows after the first match (row +/// 5 │ 400 │ ─┘ 2 buffered side; 100 < 200) +/// └──────────────────┘ +/// +/// Processing Row 2: +/// By sorting the streamed side we know +/// +/// Sorted Buffered Side Sorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 100 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ <- Start here when probing for the 2 │ 200 │ +/// ├──────────────────┤ streamed side row 2. ├──────────────────┤ +/// 3 │ 200 │ 3 │ 500 │ +/// ├──────────────────┤ └──────────────────┘ +/// 4 │ 300 │ +/// ├──────────────────┤ +/// 5 │ 400 │ +/// └──────────────────┘ +/// ``` +/// +/// ## Existence Joins (Semi, Anti, Mark) +/// Existence joins are made magnitudes of times faster with a `PiecewiseMergeJoin` as we only need to find +/// the min/max value of the streamed side to be able to emit all matches on the buffered side. By putting +/// the side we need to mark onto the sorted buffer side, we can emit all these matches at once. +/// +/// For less than operations (`<`) both inputs are to be sorted in descending order and vice versa for greater +/// than (`>`) operations. `SortExec` is used to enforce sorting on the buffered side and streamed side does not +/// need to be sorted due to only needing to find the min/max. +/// +/// For Left Semi, Anti, and Mark joins we swap the inputs so that the marked side is on the buffered side. +/// +/// The pseudocode for the algorithm looks like this: +/// +/// ```text +/// // Using the example of a less than `<` operation +/// let max = max_batch(streamed_batch) +/// +/// for buffer_row in buffer_batch: +/// if buffer_row < max: +/// output buffer_batch[buffer_row:] +/// ``` +/// +/// Only need to find the min/max value and iterate through the buffered side once. +/// +/// Here is an example: +/// We perform a `JoinType::LeftSemi` with these two batches and the operator being `Operator::Lt`(<). Because +/// the operator is `Operator::Lt` we can find the minimum value in the streamed side; in this case it is 200. +/// We can then advance a pointer from the start of the buffer side until we find the first value that satisfies +/// the predicate. All rows after that first matched value satisfy the condition 200 < x so we can mark all of +/// those rows as matched. +/// +/// ```text +/// SQL statement: +/// SELECT * +/// FROM (VALUES (500), (200), (300)) AS streamed(a) +/// LEFT SEMI JOIN (VALUES (100), (200), (200), (300), (400)) AS buffered(b) +/// ON streamed.a < buffered.b; +/// +/// Sorted Buffered Side Unsorted Streamed Side +/// ┌──────────────────┐ ┌──────────────────┐ +/// 1 │ 100 │ 1 │ 500 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 2 │ 200 │ 2 │ 200 │ +/// ├──────────────────┤ ├──────────────────┤ +/// 3 │ 200 │ 3 │ 300 │ +/// ├──────────────────┤ └──────────────────┘ +/// 4 │ 300 │ ─┐ +/// ├──────────────────┤ | We emit matches for row 4 - 5 +/// 5 │ 400 │ ─┘ on the buffered side. +/// └──────────────────┘ +/// min value: 200 +/// ``` +/// +/// For both types of joins, the buffered side must be sorted ascending for `Operator::Lt` (<) or +/// `Operator::LtEq` (<=) and descending for `Operator::Gt` (>) or `Operator::GtEq` (>=). +/// +/// # Partitioning Logic +/// Piecewise Merge Join requires one buffered side partition + round robin partitioned stream side. A counter +/// is used in the buffered side to coordinate when all streamed partitions are finished execution. This allows +/// for processing the rest of the unmatched rows for Left and Full joins. The last partition that finishes +/// execution will be responsible for outputting the unmatched rows. +/// +/// # Performance Explanation (cost) +/// Piecewise Merge Join is used over Nested Loop Join due to its superior performance. Here is the breakdown: +/// +/// R: Buffered Side +/// S: Streamed Side +/// +/// ## Piecewise Merge Join (PWMJ) +/// +/// # Classic Join: +/// Requires sorting the probe side and, for each probe row, scanning the buffered side until the first match +/// is found. +/// Complexity: `O(sort(S) + num_of_batches(|S|) * scan(R))`. +/// +/// # Mark Join: +/// Sorts the probe side, then computes the min/max range of the probe keys and scans the buffered side only +/// within that range. +/// Complexity: `O(|S| + scan(R[range]))`. +/// +/// ## Nested Loop Join +/// Compares every row from `S` with every row from `R`. +/// Complexity: `O(|S| * |R|)`. +/// +/// ## Nested Loop Join +/// Always going to be probe (O(S) * O(R)). +/// +/// # Further Reference Material +/// DuckDB blog on Range Joins: [Range Joins in DuckDB](https://duckdb.org/2022/05/27/iejoin.html) +#[derive(Debug)] +pub struct PiecewiseMergeJoinExec { + /// Left buffered execution plan + pub buffered: Arc<dyn ExecutionPlan>, + /// Right streamed execution plan + pub streamed: Arc<dyn ExecutionPlan>, + /// The two expressions being compared + pub on: (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>), + /// Comparison operator in the range predicate + pub operator: Operator, + /// How the join is performed + pub join_type: JoinType, + /// The schema once the join is applied + schema: SchemaRef, + /// Buffered data + buffered_fut: OnceAsync<BufferedSideData>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + + /// Sort expressions - See above for more details [`PiecewiseMergeJoinExec`] + /// + /// The left sort order, descending for `<`, `<=` operations + ascending for `>`, `>=` operations + left_child_plan_required_order: LexOrdering, + /// The right sort order, descending for `<`, `<=` operations + ascending for `>`, `>=` operations + /// Unsorted for mark joins + right_batch_required_orders: LexOrdering, + + /// This determines the sort order of all join columns used in sorting the stream and buffered execution plans. + sort_options: SortOptions, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// Number of partitions to process + num_partitions: usize, +} + +impl PiecewiseMergeJoinExec { + pub fn try_new( + buffered: Arc<dyn ExecutionPlan>, + streamed: Arc<dyn ExecutionPlan>, + on: (Arc<dyn PhysicalExpr>, Arc<dyn PhysicalExpr>), + operator: Operator, + join_type: JoinType, + num_partitions: usize, + ) -> Result<Self> { + // TODO: Implement existence joins for PiecewiseMergeJoin + if is_existence_join(join_type) { + return not_impl_err!( + "Existence Joins are currently not supported for PiecewiseMergeJoin" + ); + } + + // Take the operator and enforce a sort order on the streamed + buffered side based on + // the operator type. + let sort_options = match operator { + Operator::Lt | Operator::LtEq => { + // For left existence joins the inputs will be swapped so the sort + // options are switched + if is_right_existence_join(join_type) { + SortOptions::new(false, true) + } else { + SortOptions::new(true, true) + } + } + Operator::Gt | Operator::GtEq => { + if is_right_existence_join(join_type) { + SortOptions::new(true, true) + } else { + SortOptions::new(false, true) + } + } + _ => { + return internal_err!( + "Cannot contain non-range operator in PiecewiseMergeJoinExec" + ); + } + }; + + // Give the same `sort_option for comparison later` + let left_child_plan_required_order = + vec![PhysicalSortExpr::new(Arc::clone(&on.0), sort_options)]; + let right_batch_required_orders = + vec![PhysicalSortExpr::new(Arc::clone(&on.1), sort_options)]; + + let Some(left_child_plan_required_order) = + LexOrdering::new(left_child_plan_required_order) + else { + return internal_err!( + "PiecewiseMergeJoinExec requires valid sort expressions for its left side" + ); + }; + let Some(right_batch_required_orders) = + LexOrdering::new(right_batch_required_orders) + else { + return internal_err!( + "PiecewiseMergeJoinExec requires valid sort expressions for its right side" + ); + }; + + let buffered_schema = buffered.schema(); + let streamed_schema = streamed.schema(); + + // Create output schema for the join + let schema = + Arc::new(build_join_schema(&buffered_schema, &streamed_schema, &join_type).0); + let cache = Self::compute_properties( + &buffered, + &streamed, + Arc::clone(&schema), + join_type, + &on, + )?; + + Ok(Self { + streamed, + buffered, + on, + operator, + join_type, + schema, + buffered_fut: Default::default(), + metrics: ExecutionPlanMetricsSet::new(), + left_child_plan_required_order, + right_batch_required_orders, + sort_options, + cache: Arc::new(cache), + num_partitions, + }) + } + + /// Reference to buffered side execution plan + pub fn buffered(&self) -> &Arc<dyn ExecutionPlan> { + &self.buffered + } + + /// Reference to streamed side execution plan + pub fn streamed(&self) -> &Arc<dyn ExecutionPlan> { + &self.streamed + } + + /// Join type + pub fn join_type(&self) -> JoinType { + self.join_type + } + + /// Reference to sort options + pub fn sort_options(&self) -> &SortOptions { + &self.sort_options + } + + /// Get probe side (streamed side) for the PiecewiseMergeJoin + /// In current implementation, probe side is determined according to join type. + pub fn probe_side(join_type: &JoinType) -> JoinSide { + match join_type { + JoinType::Right + | JoinType::Inner + | JoinType::Full + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark => JoinSide::Right, + JoinType::Left + | JoinType::LeftAnti + | JoinType::LeftSemi + | JoinType::LeftMark => JoinSide::Left, + } + } + + pub fn compute_properties( + buffered: &Arc<dyn ExecutionPlan>, + streamed: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + join_type: JoinType, + join_on: &(PhysicalExprRef, PhysicalExprRef), + ) -> Result<PlanProperties> { + let eq_properties = join_equivalence_properties( + buffered.equivalence_properties().clone(), + streamed.equivalence_properties().clone(), + &join_type, + schema, + &Self::maintains_input_order(join_type), + Some(Self::probe_side(&join_type)), + std::slice::from_ref(join_on), + )?; + + let output_partitioning = + asymmetric_join_output_partitioning(buffered, streamed, &join_type)?; + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + boundedness_from_children([buffered, streamed]), + )) + } + + // TODO: Add input order. Now they're all `false` indicating it will not maintain the input order. + // However, for certain join types the order is maintained. This can be updated in the future after + // more testing. + fn maintains_input_order(join_type: JoinType) -> Vec<bool> { + match join_type { + // The existence side is expected to come in sorted + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + vec![false, false] + } + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + vec![false, false] + } + // Left, Right, Full, Inner Join is not guaranteed to maintain + // input order as the streamed side will be sorted during + // execution for `PiecewiseMergeJoin` + _ => vec![false, false], + } + } + + // TODO + pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> { + todo!() + } +} + +impl ExecutionPlan for PiecewiseMergeJoinExec { + fn name(&self) -> &str { + "PiecewiseMergeJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.buffered, &self.streamed] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + // Apply to the two expressions being compared in the range predicate + crate::apply_expression_roots([&self.on.0, &self.on.1], f) + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::UnspecifiedDistribution, + ]) + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + // Existence joins don't need to be sorted on one side. + if is_right_existence_join(self.join_type) { + unimplemented!() + } else { + // Sort the right side in memory, so we do not need to enforce any sorting + vec![ + Some(OrderingRequirements::from( + self.left_child_plan_required_order.clone(), + )), + None, + ] + } + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let buffered = children.swap_remove(0); + let streamed = children.swap_remove(0); + Ok(Arc::new(Self { + buffered, + streamed, + on: self.on.clone(), + operator: self.operator, + join_type: self.join_type, + schema: Arc::clone(&self.schema), + left_child_plan_required_order: self + .left_child_plan_required_order + .clone(), + right_batch_required_orders: self.right_batch_required_orders.clone(), + sort_options: self.sort_options, + cache: Arc::clone(&self.cache), + num_partitions: self.num_partitions, + + // Re-set state. + metrics: ExecutionPlanMetricsSet::new(), + buffered_fut: Default::default(), + })) + } + ChildrenPropertiesMode::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.operator, + self.join_type, + self.num_partitions, + )?)), + _ => internal_err!( + "PiecewiseMergeJoin should have 2 children, found {}", + children.len() + ), + }, + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + let buffered = Arc::clone(&self.buffered); + let streamed = Arc::clone(&self.streamed); + self.replace_children( + vec![buffered, streamed], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<datafusion_execution::TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let on_buffered = Arc::clone(&self.on.0); + let on_streamed = Arc::clone(&self.on.1); + + let metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); + let buffered_fut = self.buffered_fut.try_once(|| { + let reservation = MemoryConsumer::new("PiecewiseMergeJoinInput") + .register(context.memory_pool()); + + let buffered_stream = self.buffered.execute(0, Arc::clone(&context))?; + Ok(build_buffered_data( + buffered_stream, + Arc::clone(&on_buffered), + metrics.clone(), + reservation, + build_visited_indices_map(self.join_type), + self.num_partitions, + )) + })?; + + let streamed = self.streamed.execute(partition, Arc::clone(&context))?; + + let batch_size = context.session_config().batch_size(); + + // TODO: Add existence joins + this is guarded at physical planner + if is_existence_join(self.join_type()) { + unreachable!() + } else { + Ok(Box::pin(ClassicPWMJStream::try_new( + Arc::clone(&self.schema), + on_streamed, + self.join_type, + self.operator, + streamed, + BufferedSide::Initial(BufferedSideInitialState { buffered_fut }), + PiecewiseMergeJoinStreamState::WaitBufferedSide, + self.sort_options, + metrics, + batch_size, + ))) + } + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +impl DisplayAs for PiecewiseMergeJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + let on_str = format!( + "({} {} {})", + fmt_sql(self.on.0.as_ref()), + self.operator, + fmt_sql(self.on.1.as_ref()) + ); + + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "PiecewiseMergeJoin: operator={:?}, join_type={:?}, on={}", + self.operator, self.join_type, on_str + ) + } + + DisplayFormatType::TreeRender => { + writeln!(f, "operator={:?}", self.operator)?; + if self.join_type != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + writeln!(f, "on={on_str}") + } + } + } +} + +async fn build_buffered_data( + buffered: SendableRecordBatchStream, + on_buffered: PhysicalExprRef, + metrics: BuildProbeJoinMetrics, + reservation: MemoryReservation, + build_map: bool, + remaining_partitions: usize, +) -> Result<BufferedSideData> { + let schema = buffered.schema(); + + // Combine batches and record number of rows + let initial = (Vec::new(), 0, metrics, reservation); + let (batches, num_rows, metrics, reservation) = buffered + .try_fold(initial, |mut acc, batch| async { + let batch_size = get_record_batch_memory_size(&batch); + acc.3.try_grow(batch_size)?; + acc.2.build_mem_used.add(batch_size); + acc.2.build_input_batches.add(1); + acc.2.build_input_rows.add(batch.num_rows()); + // Update row count + acc.1 += batch.num_rows(); + // Push batch to output + acc.0.push(batch); + Ok(acc) + }) + .await?; + + let single_batch = concat_batches(&schema, batches.iter())?; + + // Evaluate physical expression on the buffered side. + let buffered_values = on_buffered + .evaluate(&single_batch)? + .into_array(single_batch.num_rows())?; + + // We add the single batch size + the memory of the join keys + // size of the size estimation + let size_estimation = get_record_batch_memory_size(&single_batch) + + buffered_values.get_array_memory_size(); + reservation.try_grow(size_estimation)?; + metrics.build_mem_used.add(size_estimation); + + // Created visited indices bitmap only if the join type requires it + let visited_indices_bitmap = if build_map { + let bitmap_size = bit_util::ceil(single_batch.num_rows(), 8); + reservation.try_grow(bitmap_size)?; + metrics.build_mem_used.add(bitmap_size); + + let mut bitmap_buffer = BooleanBufferBuilder::new(single_batch.num_rows()); + bitmap_buffer.append_n(num_rows, false); + bitmap_buffer + } else { + BooleanBufferBuilder::new(0) + }; + + let buffered_data = BufferedSideData::new( + single_batch, + buffered_values, + Mutex::new(visited_indices_bitmap), + remaining_partitions, + reservation, + ); + + Ok(buffered_data) +} + +pub(super) struct BufferedSideData { + pub(super) batch: RecordBatch, + values: ArrayRef, + pub(super) visited_indices_bitmap: SharedBitmapBuilder, + pub(super) remaining_partitions: AtomicUsize, + _reservation: MemoryReservation, +} + +impl BufferedSideData { + pub(super) fn new( + batch: RecordBatch, + values: ArrayRef, + visited_indices_bitmap: SharedBitmapBuilder, + remaining_partitions: usize, + reservation: MemoryReservation, + ) -> Self { + Self { + batch, + values, + visited_indices_bitmap, + remaining_partitions: AtomicUsize::new(remaining_partitions), + _reservation: reservation, + } + } + + pub(super) fn batch(&self) -> &RecordBatch { + &self.batch + } + + pub(super) fn values(&self) -> &ArrayRef { + &self.values + } +} + +pub(super) enum BufferedSide { + /// Indicates that build-side not collected yet + Initial(BufferedSideInitialState), + /// Indicates that build-side data has been collected + Ready(BufferedSideReadyState), +} + +impl BufferedSide { + // Takes a mutable state of the buffered row batches + pub(super) fn try_as_initial_mut(&mut self) -> Result<&mut BufferedSideInitialState> { + match self { + BufferedSide::Initial(state) => Ok(state), + _ => internal_err!("Expected build side in initial state"), + } + } + + pub(super) fn try_as_ready(&self) -> Result<&BufferedSideReadyState> { + match self { + BufferedSide::Ready(state) => Ok(state), + _ => { + internal_err!("Expected build side in ready state") + } + } + } + + /// Tries to extract BuildSideReadyState from BuildSide enum. + /// Returns an error if state is not Ready. + pub(super) fn try_as_ready_mut(&mut self) -> Result<&mut BufferedSideReadyState> { + match self { + BufferedSide::Ready(state) => Ok(state), + _ => internal_err!("Expected build side in ready state"), + } + } +} + +pub(super) struct BufferedSideInitialState { + pub(crate) buffered_fut: OnceFut<BufferedSideData>, +} + +pub(super) struct BufferedSideReadyState { + /// Collected build-side data + pub(super) buffered_data: Arc<BufferedSideData>, +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/mod.rs b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/mod.rs new file mode 100644 index 0000000..c85a7cc --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/mod.rs
@@ -0,0 +1,24 @@ +// 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 is currently experimental + +pub use exec::PiecewiseMergeJoinExec; + +mod classic_join; +mod exec; +mod utils;
diff --git a/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/utils.rs b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/utils.rs new file mode 100644 index 0000000..5bbb496 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/piecewise_merge_join/utils.rs
@@ -0,0 +1,61 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion_expr::JoinType; + +// Returns boolean for whether the join is a right existence join +pub(super) fn is_right_existence_join(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::RightAnti | JoinType::RightSemi | JoinType::RightMark + ) +} + +// Returns boolean for whether the join is an existence join +pub(super) fn is_existence_join(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftMark + | JoinType::RightMark + ) +} + +// Returns boolean to check if the join type needs to record +// buffered side matches for classic joins +pub(super) fn need_produce_result_in_final(join_type: JoinType) -> bool { + matches!(join_type, JoinType::Full | JoinType::Left) +} + +// Returns boolean for whether or not we need to build the buffered side +// bitmap for marking matched rows on the buffered side. +pub(super) fn build_visited_indices_map(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::Full + | JoinType::Left + | JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftMark + | JoinType::RightMark + ) +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/proto.rs b/native/vendor/datafusion-physical-plan/src/joins/proto.rs new file mode 100644 index 0000000..2272828 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/proto.rs
@@ -0,0 +1,161 @@ +// 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. + +//! Protobuf conversions shared by the join operators' `try_to_proto` / +//! `try_from_proto` implementations. +//! +//! The enum conversions are by-name exhaustive matches on purpose: the proto +//! enums and the `datafusion_common` enums are numbered differently, so a +//! numeric cast would silently corrupt them. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{ + JoinSide, JoinType, NullEquality, Result, internal_datafusion_err, +}; +use datafusion_proto_models::protobuf; + +use crate::joins::utils::{ColumnIndex, JoinFilter}; +use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; + +pub(crate) fn join_type_to_proto(join_type: JoinType) -> protobuf::JoinType { + match join_type { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + } +} + +pub(crate) fn join_type_from_proto(value: i32, plan_name: &str) -> Result<JoinType> { + let join_type = protobuf::JoinType::try_from(value) + .map_err(|_| internal_datafusion_err!("{plan_name}: unknown JoinType {value}"))?; + Ok(match join_type { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }) +} + +pub(crate) fn join_side_to_proto(side: JoinSide) -> protobuf::JoinSide { + match side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + } +} + +pub(crate) fn join_side_from_proto(value: i32, plan_name: &str) -> Result<JoinSide> { + let side = protobuf::JoinSide::try_from(value) + .map_err(|_| internal_datafusion_err!("{plan_name}: unknown JoinSide {value}"))?; + Ok(match side { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }) +} + +pub(crate) fn null_equality_to_proto( + null_equality: NullEquality, +) -> protobuf::NullEquality { + match null_equality { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + } +} + +pub(crate) fn null_equality_from_proto( + value: i32, + plan_name: &str, +) -> Result<NullEquality> { + let null_equality = protobuf::NullEquality::try_from(value).map_err(|_| { + internal_datafusion_err!("{plan_name}: unknown NullEquality {value}") + })?; + Ok(match null_equality { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }) +} + +pub(crate) fn join_filter_to_proto( + filter: &JoinFilter, + ctx: &ExecutionPlanEncodeCtx<'_>, +) -> Result<protobuf::JoinFilter> { + let expression = ctx.encode_expr(filter.expression())?; + let column_indices = filter + .column_indices() + .iter() + .map(|column_index| protobuf::ColumnIndex { + index: column_index.index as u32, + side: join_side_to_proto(column_index.side).into(), + }) + .collect(); + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(filter.schema().as_ref().try_into()?), + }) +} + +pub(crate) fn join_filter_from_proto( + filter: &protobuf::JoinFilter, + ctx: &ExecutionPlanDecodeCtx<'_>, + plan_name: &str, +) -> Result<JoinFilter> { + let schema: Schema = filter + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!("{plan_name}: JoinFilter missing schema") + })? + .try_into()?; + let expression = ctx.decode_required_expr( + filter.expression.as_ref(), + &schema, + plan_name, + "filter.expression", + )?; + let column_indices = filter + .column_indices + .iter() + .map(|column_index| { + Ok(ColumnIndex { + index: column_index.index as usize, + side: join_side_from_proto(column_index.side, plan_name)?, + }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/bitwise_stream.rs new file mode 100644 index 0000000..1b90f24 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/bitwise_stream.rs
@@ -0,0 +1,1265 @@ +// 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 stream specialized for semi/anti/mark joins. +//! +//! Instantiated by [`SortMergeJoinExec`](crate::joins::sort_merge_join::SortMergeJoinExec) +//! when the join type is `LeftSemi`, `LeftAnti`, `RightSemi`, `RightAnti`, +//! `LeftMark`, or `RightMark`. +//! +//! # Motivation +//! +//! The general-purpose `MaterializingSortMergeJoinStream` +//! handles semi/anti joins by materializing `(outer, inner)` row pairs, +//! applying a filter, then using a "corrected filter mask" to deduplicate. +//! Semi/anti joins only need a boolean per outer row (does a match exist?), +//! not pairs. The pair-based approach incurs unnecessary memory allocation +//! and intermediate batches. +//! +//! This stream instead tracks matches with a per-outer-batch bitset, +//! avoiding all pair materialization. +//! +//! # "Outer Side" vs "Inner Side" +//! +//! For `Left*` join types, left is outer and right is inner. +//! For `Right*` join types, right is outer and left is inner. +//! The output schema always equals the outer side's schema (for semi/anti) +//! or the outer side's schema plus a boolean mark column (for mark joins). +//! +//! # Algorithm +//! +//! Both inputs must be sorted by the join keys. The stream performs a merge +//! scan across the two sorted inputs: +//! +//! ```text +//! outer cursor ──► [1, 2, 2, 3, 5, 5, 7] +//! inner cursor ──► [2, 2, 4, 5, 6, 7, 7] +//! ▲ +//! compare keys at cursors +//! ``` +//! +//! At each step, the keys at the outer and inner cursors are compared: +//! +//! - **outer < inner**: Skip the outer key group (no match exists). +//! - **outer > inner**: Skip the inner key group. +//! - **outer == inner**: Process the match (see below). +//! +//! Key groups are contiguous runs of equal keys within one side. The scan +//! advances past entire groups at each step. +//! +//! ## Processing a key match +//! +//! **Without filter**: All outer rows in the key group are marked as matched. +//! +//! **With filter**: The inner key group is buffered (may span multiple inner +//! batches). For each buffered inner row, the filter is evaluated against the +//! outer key group as a batch. Results are OR'd into the matched bitset. A +//! short-circuit exits early when all outer rows in the group are matched. +//! +//! ```text +//! matched bitset: [0, 0, 1, 0, 0, ...] +//! ▲── one bit per outer row ──▲ +//! +//! On emit: +//! Semi → filter_record_batch(outer_batch, &matched) +//! Anti → filter_record_batch(outer_batch, &NOT(matched)) +//! Mark → outer_batch + matched as boolean column +//! ``` +//! +//! ## Batch boundaries +//! +//! Key groups can span batch boundaries on either side. The stream handles +//! this by detecting when a group extends to the end of a batch, loading the +//! next batch, and continuing if the key matches. The generator-based stream +//! suspends in place at `await` points, so no explicit re-entry state is +//! needed. +//! +//! # Memory +//! +//! Memory usage is bounded and independent of total input size: +//! - One outer batch at a time (not tracked by reservation — single batch, +//! cannot be spilled since it's needed for filter evaluation) +//! - One inner batch at a time (streaming) +//! - `matched` bitset: one bit per outer row, re-allocated per batch +//! - Inner key group buffer: only for filtered joins, one key group at a time. +//! Tracked via `MemoryReservation`; spilled to disk when the memory pool +//! limit is exceeded. +//! - `BatchCoalescer`: output buffering to target batch size +//! +//! # Degenerate cases +//! +//! **Highly skewed key (filtered joins only):** When a filter is present, +//! the inner key group is buffered so each inner row can be evaluated +//! against the outer group. If one join key has N inner rows, all N rows +//! are held in memory simultaneously (or spilled to disk if the memory +//! pool limit is reached). With uniform key distribution this is small +//! (inner_rows / num_distinct_keys), but a single hot key can buffer +//! arbitrarily many rows. The no-filter path does not buffer inner +//! rows — it only advances the cursor — so it is unaffected. +//! +//! **Scalar broadcast during filter evaluation:** Each inner row is +//! broadcast to match the outer group length for filter evaluation, +//! allocating one array per inner row × filter column. This is inherent +//! to the `PhysicalExpr::evaluate(RecordBatch)` API, which does not +//! support scalar inputs directly. The total work is +//! O(inner_group × outer_group) per key, but with much lower constant +//! factor than the pair-materialization approach. + +use std::cmp::Ordering; +use std::sync::Arc; + +use crate::EmptyRecordBatchStream; +use crate::joins::utils::{JoinFilter, JoinKeyComparator, compare_join_arrays}; +use crate::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, Time, +}; +use crate::spill::in_progress_spill_file::InProgressSpillFile; +use crate::spill::spill_manager::SpillManager; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; +use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch}; +use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not}; +use arrow::datatypes::SchemaRef; +use arrow::util::bit_chunk_iterator::UnalignedBitChunk; +use arrow::util::bit_util::apply_bitwise_binary_op; +use datafusion_common::instant::Instant; +use datafusion_common::{ + DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err, +}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{ + SendableRecordBatchStream, SpillFile, TryEmitter, async_try_stream, +}; +use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; + +use futures::StreamExt; + +/// Evaluates join key expressions against a batch, returning one array per key. +fn evaluate_join_keys( + batch: &RecordBatch, + on: &[PhysicalExprRef], +) -> Result<Vec<ArrayRef>> { + on.iter() + .map(|expr| { + let num_rows = batch.num_rows(); + let val = expr.evaluate(batch)?; + val.into_array(num_rows) + }) + .collect() +} + +/// Find the first index in `key_arrays` starting from `from` where the key +/// differs from the key at `from`. Uses a pre-built `JoinKeyComparator` for +/// zero-alloc ordinal comparison without per-row type dispatch. +/// +/// Optimized for join workloads: checks adjacent and boundary keys before +/// falling back to binary search, since most key groups are small (often 1). +fn find_key_group_end(cmp: &JoinKeyComparator, from: usize, len: usize) -> usize { + let next = from + 1; + if next >= len { + return len; + } + + // Fast path: single-row group (common with unique keys). + if cmp.compare(from, next) != Ordering::Equal { + return next; + } + + // Check if the entire remaining batch shares this key. + let last = len - 1; + if cmp.compare(from, last) == Ordering::Equal { + return len; + } + + // Binary search the interior: key at `next` matches, key at `last` doesn't. + let mut lo = next + 1; + let mut hi = last; + while lo < hi { + let mid = lo + (hi - lo) / 2; + if cmp.compare(from, mid) == Ordering::Equal { + lo = mid + 1; + } else { + hi = mid; + } + } + lo +} + +/// Sort-Merge join stream for Semi/Anti/Mark joins. +/// +/// Named "bitwise" because it tracks outer-row matches via a per-batch +/// boolean bitset (`BooleanBufferBuilder`) rather than materializing +/// `(outer, inner)` row pairs. Filter results are OR'd into the bitset +/// in `u64` chunks, and emitting applies the bitset directly. +pub(crate) struct BitwiseSortMergeJoinStream { + join_type: JoinType, + + // Input streams — in the nested-loop model that sort-merge join + // implements, "outer" is the driving loop and "inner" is probed for + // matches. The existing MaterializingSortMergeJoinStream calls these "streamed" + // and "buffered" respectively. For Left* joins, outer=left; for + // Right* joins, outer=right. Output schema equals the outer side. + outer: SendableRecordBatchStream, + inner: SendableRecordBatchStream, + + // Current batches and cursor positions within them + outer_batch: Option<RecordBatch>, + /// Row index into `outer_batch` — the next unprocessed outer row. + outer_offset: usize, + outer_key_arrays: Vec<ArrayRef>, + inner_batch: Option<RecordBatch>, + /// Row index into `inner_batch` — the next unprocessed inner row. + inner_offset: usize, + inner_key_arrays: Vec<ArrayRef>, + + // Per-outer-batch match tracking, reused across batches. + // Bit-packed (not Vec<bool>) so that: + // - emit: finish() yields a BooleanBuffer directly (no packing iteration) + // - OR: apply_bitwise_binary_op ORs filter results in u64 chunks + // - count: UnalignedBitChunk::count_ones uses popcnt + matched: BooleanBufferBuilder, + + // Inner key group buffer: all inner rows sharing the current join key. + // Only populated when a filter is present. Unbounded — a single key + // with many inner rows will buffer them all. See "Degenerate cases" + // in exec.rs. On memory pool overflow the buffered slices move to a + // per-group spill file (see [`Self::buffer_inner_key_group`]). + inner_key_buffer: Vec<RecordBatch>, + + // Join ON expressions, evaluated against each new batch to produce + // the key arrays used for sorted key comparisons. + on_outer: Vec<PhysicalExprRef>, + on_inner: Vec<PhysicalExprRef>, + filter: Option<JoinFilter>, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, + // Decomposed from JoinType: when RightSemi/RightAnti, outer=right, + // inner=left, so we swap sides when building the filter batch. + outer_is_left: bool, + + // Output + coalescer: BatchCoalescer, + schema: SchemaRef, + + // Metrics — output rows/batches and end time are recorded by the + // ObservedStream wrapper in try_new, not here. + input_batches: Count, + input_rows: Count, + peak_mem_used: Gauge, + /// Time spent doing the join's own work (including spill write and + /// read-back). The clock is stopped while awaiting the child inputs or + /// the consumer taking an emitted batch — see [`Self::stop_join_time`]. + join_time: Time, + /// Start of the currently running `join_time` span; `None` while the + /// clock is stopped. + join_time_start: Option<Instant>, + + // Memory / spill — only the inner key buffer is tracked via reservation, + // matching existing SMJ (which tracks only the buffered side). The outer + // batch is a single batch at a time and cannot be spilled. + reservation: MemoryReservation, + spill_manager: SpillManager, + runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>, + inner_buffer_size: usize, + + // Cached comparators — pre-built to avoid per-row type dispatch. + /// Comparator for outer vs inner key comparison + outer_inner_cmp: Option<JoinKeyComparator>, + /// Comparator for outer self-comparison (find_key_group_end on outer) + outer_self_cmp: Option<JoinKeyComparator>, + /// Comparator for inner self-comparison (find_key_group_end on inner) + inner_self_cmp: Option<JoinKeyComparator>, +} + +impl BitwiseSortMergeJoinStream { + #[expect(clippy::too_many_arguments)] + pub fn try_new( + schema: SchemaRef, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, + outer: SendableRecordBatchStream, + inner: SendableRecordBatchStream, + on_outer: Vec<PhysicalExprRef>, + on_inner: Vec<PhysicalExprRef>, + filter: Option<JoinFilter>, + join_type: JoinType, + batch_size: usize, + partition: usize, + metrics: &ExecutionPlanMetricsSet, + reservation: MemoryReservation, + spill_manager: SpillManager, + runtime_env: Arc<datafusion_execution::runtime_env::RuntimeEnv>, + ) -> Result<SendableRecordBatchStream> { + debug_assert!( + matches!( + join_type, + JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ), + "BitwiseSortMergeJoinStream does not handle {join_type:?}" + ); + let outer_is_left = matches!( + join_type, + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark + ); + + let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition); + let input_batches = + MetricBuilder::new(metrics).counter("input_batches", partition); + let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition); + let baseline_metrics = BaselineMetrics::new(metrics, partition); + let peak_mem_used = + MetricBuilder::new(metrics).peak_memory_usage("peak_mem_used", partition); + + let mut state = Self { + join_type, + outer, + inner, + outer_batch: None, + outer_offset: 0, + outer_key_arrays: vec![], + inner_batch: None, + inner_offset: 0, + inner_key_arrays: vec![], + matched: BooleanBufferBuilder::new(0), + inner_key_buffer: vec![], + on_outer, + on_inner, + filter, + sort_options, + null_equality, + outer_is_left, + coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size) + .with_biggest_coalesce_batch_size(Some(batch_size / 2)), + schema: Arc::clone(&schema), + input_batches, + input_rows, + peak_mem_used, + join_time, + join_time_start: None, + reservation, + spill_manager, + runtime_env, + inner_buffer_size: 0, + outer_inner_cmp: None, + outer_self_cmp: None, + inner_self_cmp: None, + }; + + let stream = async_try_stream(|mut emitter| async move { + state.start_join_time(); + let result = state.join(&mut emitter).await; + state.stop_join_time(); + result + }); + // ObservedStream records the baseline metrics (output rows/batches, + // end time) exactly as the former hand-written poll_next did. + Ok(Box::pin(ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, + None, + ))) + } + + /// Start (resume) the `join_time` clock. + fn start_join_time(&mut self) { + debug_assert!(self.join_time_start.is_none(), "join_time already running"); + self.join_time_start = Some(Instant::now()); + } + + /// Stop (pause) the `join_time` clock, accumulating the elapsed span. + /// + /// Called around awaits whose duration is not the join's own work: the + /// child input streams' `next()` and `emitter.emit()` (where the + /// consumer processes the batch). The join's own spill read-back is NOT + /// excluded — that time is join work. + fn stop_join_time(&mut self) { + if let Some(start) = self.join_time_start.take() { + self.join_time.add_elapsed(start); + } + } + + /// Resize the memory reservation to match current tracked usage. + fn try_resize_reservation(&mut self) -> Result<()> { + let needed = self.inner_buffer_size; + self.reservation.try_resize(needed)?; + self.peak_mem_used.set_max(self.reservation.size()); + Ok(()) + } + + /// Get or build the outer vs inner key comparator. + fn get_outer_inner_cmp(&mut self) -> Result<&JoinKeyComparator> { + if self.outer_inner_cmp.is_none() { + self.outer_inner_cmp = Some(JoinKeyComparator::new( + &self.outer_key_arrays, + &self.inner_key_arrays, + &self.sort_options, + self.null_equality, + )?); + } + Ok(self.outer_inner_cmp.as_ref().unwrap()) + } + + /// Get or build the outer self-comparison comparator. + fn get_outer_self_cmp(&mut self) -> Result<&JoinKeyComparator> { + if self.outer_self_cmp.is_none() { + self.outer_self_cmp = Some(JoinKeyComparator::new( + &self.outer_key_arrays, + &self.outer_key_arrays, + &self.sort_options, + self.null_equality, + )?); + } + Ok(self.outer_self_cmp.as_ref().unwrap()) + } + + /// Get or build the inner self-comparison comparator. + fn get_inner_self_cmp(&mut self) -> Result<&JoinKeyComparator> { + if self.inner_self_cmp.is_none() { + self.inner_self_cmp = Some(JoinKeyComparator::new( + &self.inner_key_arrays, + &self.inner_key_arrays, + &self.sort_options, + self.null_equality, + )?); + } + Ok(self.inner_self_cmp.as_ref().unwrap()) + } + + /// Spill the in-memory inner key buffer to disk and clear it. One key + /// group can spill repeatedly; every call appends to `writer` — the + /// group's single open spill file — creating it on first use. + fn spill_inner_key_buffer( + &mut self, + writer: &mut Option<InProgressSpillFile>, + ) -> Result<()> { + if writer.is_none() { + *writer = Some( + self.spill_manager + .create_in_progress_file("semi_anti_smj_inner_key_spill")?, + ); + } + let writer = writer.as_mut().unwrap(); + for batch in self.inner_key_buffer.drain(..) { + writer.append_batch(&batch)?; + } + self.inner_buffer_size = 0; + // Should succeed now — inner buffer has been spilled. + self.try_resize_reservation() + } + + /// Clear inner key group state after processing. Does not resize the + /// reservation — the next key group will resize when buffering, or + /// the stream's Drop will free it. This avoids unnecessary memory + /// pool interactions (see apache/datafusion#20729). + fn clear_inner_key_group(&mut self) { + self.inner_key_buffer.clear(); + self.inner_buffer_size = 0; + } + + /// Fetch the next outer batch. Returns true if a batch was loaded. + async fn next_outer_batch(&mut self) -> Result<bool> { + loop { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.outer.next().await; + self.start_join_time(); + match item { + None => { + // Release the outer input pipeline's resources. + let outer_schema = self.outer.schema(); + self.outer = Box::pin(EmptyRecordBatchStream::new(outer_schema)); + return Ok(false); + } + Some(Err(e)) => return Err(e), + Some(Ok(batch)) => { + let batch_num_rows = batch.num_rows(); + self.input_batches.add(1); + self.input_rows.add(batch_num_rows); + if batch_num_rows == 0 { + continue; + } + let keys = evaluate_join_keys(&batch, &self.on_outer)?; + self.outer_batch = Some(batch); + self.outer_offset = 0; + self.outer_key_arrays = keys; + self.outer_inner_cmp = None; + self.outer_self_cmp = None; + self.matched = BooleanBufferBuilder::new(batch_num_rows); + self.matched.append_n(batch_num_rows, false); + return Ok(true); + } + } + } + } + + /// Fetch the next inner batch. Returns true if a batch was loaded. + async fn next_inner_batch(&mut self) -> Result<bool> { + loop { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.inner.next().await; + self.start_join_time(); + match item { + None => { + // Release the inner input pipeline's resources. + let inner_schema = self.inner.schema(); + self.inner = Box::pin(EmptyRecordBatchStream::new(inner_schema)); + return Ok(false); + } + Some(Err(e)) => return Err(e), + Some(Ok(batch)) => { + let batch_num_rows = batch.num_rows(); + self.input_batches.add(1); + self.input_rows.add(batch_num_rows); + if batch_num_rows == 0 { + continue; + } + let keys = evaluate_join_keys(&batch, &self.on_inner)?; + self.inner_batch = Some(batch); + self.inner_offset = 0; + self.inner_key_arrays = keys; + self.outer_inner_cmp = None; + self.inner_self_cmp = None; + return Ok(true); + } + } + } + } + + /// Push the current outer batch into the coalescer, applying the matched + /// bitset as a selection mask. Consumes the batch (`outer_batch` becomes + /// `None`). + fn emit_outer_batch(&mut self) -> Result<()> { + let batch = self.outer_batch.take().unwrap(); + + // finish() converts the bit-packed builder directly to a + // BooleanBuffer — no iteration or repacking needed. + let matched_buf = self.matched.finish(); + + match self.join_type { + JoinType::LeftMark | JoinType::RightMark => { + // Mark joins emit ALL outer rows with a boolean match column appended. + debug_assert_eq!( + self.schema.fields().len(), + batch.num_columns() + 1, + "Mark join output schema should be outer schema + 1 mark column" + ); + let mark_col = Arc::new(BooleanArray::new(matched_buf, None)) as ArrayRef; + let mut columns = Vec::with_capacity(batch.num_columns() + 1); + columns.extend_from_slice(batch.columns()); + columns.push(mark_col); + let output = RecordBatch::try_new(Arc::clone(&self.schema), columns)?; + self.coalescer.push_batch(output)?; + } + JoinType::LeftSemi | JoinType::RightSemi => { + let selection = BooleanArray::new(matched_buf, None); + let filtered = filter_record_batch(&batch, &selection)?; + if filtered.num_rows() > 0 { + self.coalescer.push_batch(filtered)?; + } + } + JoinType::LeftAnti | JoinType::RightAnti => { + let selection = not(&BooleanArray::new(matched_buf, None))?; + let filtered = filter_record_batch(&batch, &selection)?; + if filtered.num_rows() > 0 { + self.coalescer.push_batch(filtered)?; + } + } + _ => unreachable!(), + } + Ok(()) + } + + /// Mark all outer rows in the current key group as matched and advance + /// the outer cursor past the group (within the current batch). + fn mark_outer_key_group_matched(&mut self) -> Result<()> { + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + let from = self.outer_offset; + let group_end = find_key_group_end(self.get_outer_self_cmp()?, from, num_outer); + + for i in from..group_end { + self.matched.set_bit(i, true); + } + + self.outer_offset = group_end; + Ok(()) + } + + /// Advance the inner cursor past the current key group. The group may + /// span multiple inner batches. Sets `inner_batch` to `None` if inner + /// is exhausted. + async fn advance_inner_past_key_group(&mut self) -> Result<()> { + loop { + let Some(inner_batch) = &self.inner_batch else { + return Ok(()); + }; + let num_inner = inner_batch.num_rows(); + let from = self.inner_offset; + let group_end = + find_key_group_end(self.get_inner_self_cmp()?, from, num_inner); + + if group_end < num_inner { + self.inner_offset = group_end; + return Ok(()); + } + + // Key group extends to the end of the batch — it may continue + // into the next one; save the last key so we can check. + let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1); + + if !self.next_inner_batch().await? { + self.inner_batch = None; + return Ok(()); + } + if !keys_match( + &saved_inner_keys, + &self.inner_key_arrays, + &self.sort_options, + self.null_equality, + )? { + return Ok(()); + } + } + } + + /// Buffer the inner key group for filter evaluation, advancing the inner + /// cursor past the group. Collects all inner rows with the current key + /// across batch boundaries. Sets `inner_batch` to `None` if inner is + /// exhausted. + /// + /// Slices that overflow the memory pool are appended to a single spill + /// file, returned finished — ready for reading — once the whole group + /// has been buffered. `None` means the group fit in memory. + async fn buffer_inner_key_group(&mut self) -> Result<Option<Arc<dyn SpillFile>>> { + self.clear_inner_key_group(); + let mut writer: Option<InProgressSpillFile> = None; + + while let Some(inner_batch) = &self.inner_batch { + let num_inner = inner_batch.num_rows(); + let from = self.inner_offset; + let group_end = + find_key_group_end(self.get_inner_self_cmp()?, from, num_inner); + + let inner_batch = self.inner_batch.as_ref().unwrap(); + let slice = inner_batch.slice(from, group_end - from); + self.inner_buffer_size += slice.get_array_memory_size(); + self.inner_key_buffer.push(slice); + + // Reserve memory for the newly buffered slice. If the pool + // is exhausted, spill the entire buffer to disk. + if self.try_resize_reservation().is_err() { + if self.runtime_env.disk_manager.tmp_files_enabled() { + self.spill_inner_key_buffer(&mut writer)?; + } else { + // Re-attempt to get the error message + self.try_resize_reservation().map_err(|e| { + DataFusionError::Execution(format!( + "{e}. Disk spilling disabled." + )) + })?; + } + } + + if group_end < num_inner { + self.inner_offset = group_end; + break; + } + + // Key group extends to the end of the batch — it may continue + // into the next one; save the last key so we can check. + let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1); + + if !self.next_inner_batch().await? { + self.inner_batch = None; + break; + } + if !keys_match( + &saved_inner_keys, + &self.inner_key_arrays, + &self.sort_options, + self.null_equality, + )? { + break; + } + } + + match writer { + Some(mut writer) => writer.finish(), + None => Ok(None), + } + } + + /// Process a key match with a filter. For each inner row in the buffered + /// key group — the spilled slices in `spill` plus the in-memory + /// `inner_key_buffer` — evaluates the filter against the outer key group + /// and ORs the results into the matched bitset using u64-chunked bitwise + /// ops. + async fn process_key_match_with_filter( + &mut self, + spill: Option<&Arc<dyn SpillFile>>, + ) -> Result<()> { + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + + // buffer_inner_key_group must be called before this function + debug_assert!( + !self.inner_key_buffer.is_empty() || spill.is_some(), + "process_key_match_with_filter called with no inner key data" + ); + debug_assert!( + self.outer_offset < num_outer, + "outer_offset must be within the current batch" + ); + debug_assert!( + self.matched.len() == num_outer, + "matched vector must be sized for the current outer batch" + ); + + let outer_group_start = self.outer_offset; + let outer_group_end = + find_key_group_end(self.get_outer_self_cmp()?, outer_group_start, num_outer); + let outer_group_len = outer_group_end - outer_group_start; + + let filter = self.filter.as_ref().unwrap(); + let outer_batch = self.outer_batch.as_ref().unwrap(); + let outer_slice = outer_batch.slice(outer_group_start, outer_group_len); + + // Count already-matched bits using popcnt on u64 chunks (zero-copy). + let mut matched_count = UnalignedBitChunk::new( + self.matched.as_slice(), + outer_group_start, + outer_group_len, + ) + .count_ones(); + + // Process spilled inner batches first asynchronously. + if matched_count < outer_group_len + && let Some(spill_file) = spill + { + let mut spill_stream = self + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; + let mut spill_stream_has_data = false; + + // Note: the clock keeps running across the spill reads — the + // spill file is the join's own data, so reading it back is + // join work (unlike the child inputs' `next()`). + while matched_count < outer_group_len { + match spill_stream.next().await { + Some(Ok(inner_slice)) => { + spill_stream_has_data = true; + matched_count = eval_filter_for_inner_slice( + self.outer_is_left, + filter, + &outer_slice, + &inner_slice, + &mut self.matched, + outer_group_start, + outer_group_len, + matched_count, + )?; + } + Some(Err(e)) => return Err(e), + None => { + if !spill_stream_has_data { + return internal_err!("Spill file was empty"); + } + break; + } + } + } + } + + // Then process in-memory inner batches. + // evaluate_filter_for_inner_row is a free function (not &self method) + // so that Rust can split the struct borrow: &mut self.matched coexists + // with &self.inner_key_buffer and &self.filter inside this loop. + if matched_count < outer_group_len { + 'outer: for inner_slice in &self.inner_key_buffer { + matched_count = eval_filter_for_inner_slice( + self.outer_is_left, + filter, + &outer_slice, + inner_slice, + &mut self.matched, + outer_group_start, + outer_group_len, + matched_count, + )?; + if matched_count == outer_group_len { + break 'outer; + } + } + } + + self.outer_offset = outer_group_end; + + Ok(()) + } + + /// Evaluate the filter for the buffered inner key group against the + /// outer key group. If the outer key group continues into subsequent + /// outer batches, keep evaluating there too. Dropping `spill` on return + /// deletes the group's temp file. + async fn process_filtered_match_loop( + &mut self, + spill: Option<Arc<dyn SpillFile>>, + ) -> Result<()> { + loop { + self.process_key_match_with_filter(spill.as_ref()).await?; + + let outer_batch = self.outer_batch.as_ref().unwrap(); + if self.outer_offset < outer_batch.num_rows() { + break; + } + + // The outer key group may continue into the next outer batch; + // save the last key so we can check. + let saved_keys = + slice_keys(&self.outer_key_arrays, outer_batch.num_rows() - 1); + + self.emit_outer_batch()?; + + if !self.next_outer_batch().await? { + break; + } + if !keys_match( + &saved_keys, + &self.outer_key_arrays, + &self.sort_options, + self.null_equality, + )? { + break; + } + } + + self.clear_inner_key_group(); + Ok(()) + } + + /// Mark the outer key group as matched. If the outer key group continues + /// into subsequent outer batches, keep marking there too. + async fn process_unfiltered_match_loop(&mut self) -> Result<()> { + loop { + self.mark_outer_key_group_matched()?; + + let outer_batch = self.outer_batch.as_ref().unwrap(); + if self.outer_offset < outer_batch.num_rows() { + return Ok(()); + } + + // The outer key group may continue into the next outer batch; + // save the last key so we can check. + let saved_keys = + slice_keys(&self.outer_key_arrays, outer_batch.num_rows() - 1); + + self.emit_outer_batch()?; + + if !self.next_outer_batch().await? { + return Ok(()); + } + if !keys_match( + &saved_keys, + &self.outer_key_arrays, + &self.sort_options, + self.null_equality, + )? { + return Ok(()); + } + } + } + + /// Keys at both cursors are equal: determine which outer rows in the key + /// group have a match. Both key groups may span batch boundaries. + async fn process_key_match(&mut self) -> Result<()> { + if self.filter.is_some() { + // Buffer the inner key group so each inner row can be evaluated + // against the outer key group, OR-ing filter results into the + // matched bitset. + let spill = self.buffer_inner_key_group().await?; + self.process_filtered_match_loop(spill).await + } else { + // Without a filter, key equality alone means every outer row in + // the group matches; the inner rows themselves are not needed. + self.advance_inner_past_key_group().await?; + self.process_unfiltered_match_loop().await + } + } + + /// Compare the join keys at the outer and inner cursors, returning the + /// ordering of the outer key relative to the inner key (e.g. `Greater` + /// means outer key > inner key, per the sort options). + fn compare_current_keys(&mut self) -> Result<Ordering> { + let (outer_idx, inner_idx) = (self.outer_offset, self.inner_offset); + Ok(self.get_outer_inner_cmp()?.compare(outer_idx, inner_idx)) + } + + /// Outer key is unmatched: advance the outer cursor past its key group + /// (within the current batch). If the group continues into the next + /// batch, those rows compare Less again and are skipped the same way. + fn skip_outer_key_group(&mut self) -> Result<()> { + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + let from = self.outer_offset; + self.outer_offset = + find_key_group_end(self.get_outer_self_cmp()?, from, num_outer); + Ok(()) + } + + /// Sync fast path for `Ordering::Greater`: skip the inner key group when + /// it ends within the current batch. Returns false — leaving all state + /// unchanged — when the group reaches the batch boundary, in which case + /// the caller must take [`Self::advance_inner_past_key_group`]. + fn try_skip_inner_key_group(&mut self) -> Result<bool> { + let num_inner = self.inner_batch.as_ref().unwrap().num_rows(); + let from = self.inner_offset; + let group_end = find_key_group_end(self.get_inner_self_cmp()?, from, num_inner); + if group_end >= num_inner { + return Ok(false); + } + self.inner_offset = group_end; + Ok(true) + } + + /// Sync fast path for `Ordering::Equal` without a filter: when both key + /// groups end within their current batches (the common case — a group + /// only reaches a batch boundary once per batch), mark the outer group + /// matched and advance both cursors without any async machinery. + /// Returns false — leaving all state unchanged — when a filter is + /// present or either group reaches a batch boundary, in which case the + /// caller must take [`Self::process_key_match`]. + fn try_process_key_match(&mut self) -> Result<bool> { + if self.filter.is_some() { + return Ok(false); + } + + let num_inner = self.inner_batch.as_ref().unwrap().num_rows(); + let inner_from = self.inner_offset; + let inner_group_end = + find_key_group_end(self.get_inner_self_cmp()?, inner_from, num_inner); + if inner_group_end >= num_inner { + return Ok(false); + } + + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + let outer_from = self.outer_offset; + let outer_group_end = + find_key_group_end(self.get_outer_self_cmp()?, outer_from, num_outer); + if outer_group_end >= num_outer { + return Ok(false); + } + + for i in outer_from..outer_group_end { + self.matched.set_bit(i, true); + } + self.outer_offset = outer_group_end; + self.inner_offset = inner_group_end; + Ok(true) + } + + /// True when the outer cursor already points at an unprocessed row: the + /// sync fast path of [`Self::advance_outer_row`]. Checked inline in the + /// hot loop so the async helper (and its state machine) is only entered + /// at batch boundaries — same pattern as `sorts/merge.rs`. + fn has_current_outer_row(&self) -> bool { + self.outer_batch + .as_ref() + .is_some_and(|batch| self.outer_offset < batch.num_rows()) + } + + /// True when the inner cursor already points at an unprocessed row: the + /// sync fast path of [`Self::advance_inner_row`]. + fn has_current_inner_row(&self) -> bool { + self.inner_batch + .as_ref() + .is_some_and(|batch| self.inner_offset < batch.num_rows()) + } + + /// Ensure the outer cursor points at an unprocessed row, emitting + /// finished outer batches and loading new ones as needed. Returns false + /// when outer is exhausted. + async fn advance_outer_row( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<bool> { + loop { + match &self.outer_batch { + Some(batch) if self.outer_offset < batch.num_rows() => { + return Ok(true); + } + Some(_) => { + // Current batch fully scanned — emit it and load the next. + self.emit_outer_batch()?; + self.emit_completed_batches(emitter).await; + } + None => { + if !self.next_outer_batch().await? { + return Ok(false); + } + } + } + } + } + + /// Ensure the inner cursor points at an unprocessed row, loading new + /// inner batches as needed. Returns false when inner is exhausted. + async fn advance_inner_row(&mut self) -> Result<bool> { + loop { + if let Some(batch) = &self.inner_batch + && self.inner_offset < batch.num_rows() + { + return Ok(true); + } + if !self.next_inner_batch().await? { + self.inner_batch = None; + return Ok(false); + } + } + } + + /// Inner is exhausted, so no further matches are possible: emit the + /// current outer batch and all remaining ones with their current matched + /// bits (semi drops unmatched rows, anti emits them, mark emits them + /// with mark=false). + async fn drain_outer(&mut self) -> Result<()> { + self.emit_outer_batch()?; + while self.next_outer_batch().await? { + self.emit_outer_batch()?; + } + Ok(()) + } + + /// Emit all completed coalescer batches to the stream consumer. + async fn emit_completed_batches( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) { + while let Some(batch) = self.coalescer.next_completed_batch() { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); + emitter.emit(batch).await; + self.start_join_time(); + } + } + + /// Main loop: a classic merge-scan over the two sorted inputs, emitting + /// output batches as they complete. + async fn join( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + // The `has_current_*` / `has_completed_batch` fast paths keep async + // state machinery out of the per-key-group hot path; the awaiting + // helpers are only entered at batch boundaries. + while self.has_current_outer_row() || self.advance_outer_row(emitter).await? { + if !(self.has_current_inner_row() || self.advance_inner_row().await?) { + self.drain_outer().await?; + break; + } + + // Each arm handles the common case synchronously (`try_*`); the + // async continuations only run when a key group reaches a batch + // boundary or a filter must be evaluated. + match self.compare_current_keys()? { + Ordering::Less => self.skip_outer_key_group()?, + Ordering::Greater => { + if !self.try_skip_inner_key_group()? { + self.advance_inner_past_key_group().await?; + } + } + Ordering::Equal => { + if !self.try_process_key_match()? { + self.process_key_match().await?; + } + } + } + + if self.coalescer.has_completed_batch() { + self.emit_completed_batches(emitter).await; + } + } + + // Flush whatever is still buffered in the coalescer. + self.coalescer.finish_buffered_batch()?; + self.emit_completed_batches(emitter).await; + Ok(()) + } +} + +/// Evaluate the filter for all rows in an inner slice against the outer group, +/// OR-ing results into the matched bitset. Returns the updated matched count. +/// Extracted as a free function so Rust can split borrows on the stream struct. +#[expect(clippy::too_many_arguments)] +fn eval_filter_for_inner_slice( + outer_is_left: bool, + filter: &JoinFilter, + outer_slice: &RecordBatch, + inner_slice: &RecordBatch, + matched: &mut BooleanBufferBuilder, + outer_offset: usize, + outer_group_len: usize, + // Passed in to avoid recounting bits we just counted at the call site. + mut matched_count: usize, +) -> Result<usize> { + debug_assert_eq!( + matched_count, + UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len) + .count_ones() + ); + for inner_row in 0..inner_slice.num_rows() { + if matched_count == outer_group_len { + break; + } + + let filter_result = evaluate_filter_for_inner_row( + outer_is_left, + filter, + outer_slice, + inner_slice, + inner_row, + )?; + + // OR filter results into the matched bitset. Both sides are + // bit-packed [u8] buffers, so apply_bitwise_binary_op + // processes 64 bits per loop iteration (not 1 bit at a time). + // + // The offsets handle alignment: outer_offset is the bit + // position within matched where this key group starts, + // and filter_buf.offset() is the BooleanBuffer's internal + // bit offset (usually 0, but not guaranteed by Arrow). + let filter_buf = filter_result.values(); + apply_bitwise_binary_op( + matched.as_slice_mut(), + outer_offset, + filter_buf.inner().as_slice(), + filter_buf.offset(), + outer_group_len, + |a, b| a | b, + ); + + // Recount matched bits after the OR. UnalignedBitChunk is + // zero-copy — it reads the bytes in place and uses popcnt. + matched_count = + UnalignedBitChunk::new(matched.as_slice(), outer_offset, outer_group_len) + .count_ones(); + } + Ok(matched_count) +} + +/// Slice each key array to a single row at `idx`. +fn slice_keys(keys: &[ArrayRef], idx: usize) -> Vec<ArrayRef> { + keys.iter().map(|a| a.slice(idx, 1)).collect() +} + +/// Compare the first row of two key arrays using sort options to determine +/// equality. The left side is expected to be single-row slices (from +/// `slice_keys`); the right side can be any length (row 0 is compared). +fn keys_match( + left_arrays: &[ArrayRef], + right_arrays: &[ArrayRef], + sort_options: &[SortOptions], + null_equality: NullEquality, +) -> Result<bool> { + debug_assert!(left_arrays.iter().all(|a| a.len() == 1)); + let cmp = compare_join_arrays( + left_arrays, + 0, + right_arrays, + 0, + sort_options, + null_equality, + )?; + Ok(cmp == Ordering::Equal) +} + +/// Evaluate the join filter for one inner row against a slice of outer rows. +/// +/// Free function (not a method on BitwiseSortMergeJoinStream) so that Rust +/// can split the struct borrow in process_key_match_with_filter: the caller +/// holds &mut self.matched and &self.inner_key_buffer simultaneously, which +/// is impossible if this borrows all of &self. +fn evaluate_filter_for_inner_row( + outer_is_left: bool, + filter: &JoinFilter, + outer_slice: &RecordBatch, + inner_batch: &RecordBatch, + inner_idx: usize, +) -> Result<BooleanArray> { + let num_outer_rows = outer_slice.num_rows(); + + // Build filter input columns in the order the filter expects + let mut columns: Vec<ArrayRef> = Vec::with_capacity(filter.column_indices().len()); + for col_idx in filter.column_indices() { + let (side_batch, side_idx) = if outer_is_left { + match col_idx.side { + JoinSide::Left => (outer_slice, None), + JoinSide::Right => (inner_batch, Some(inner_idx)), + JoinSide::None => { + return internal_err!("Unexpected JoinSide::None in filter"); + } + } + } else { + match col_idx.side { + JoinSide::Left => (inner_batch, Some(inner_idx)), + JoinSide::Right => (outer_slice, None), + JoinSide::None => { + return internal_err!("Unexpected JoinSide::None in filter"); + } + } + }; + + match side_idx { + None => { + columns.push(Arc::clone(side_batch.column(col_idx.index))); + } + Some(idx) => { + // Broadcasts inner scalar to N-element array. Arrow's + // BinaryExpr handles Scalar×Array natively via the Datum + // trait, but Column::evaluate always returns Array, so + // we'd need a custom expr to avoid this broadcast. + let scalar = ScalarValue::try_from_array( + side_batch.column(col_idx.index).as_ref(), + idx, + )?; + columns.push(scalar.to_array_of_size(num_outer_rows)?); + } + } + } + + let filter_batch = RecordBatch::try_new(Arc::clone(filter.schema()), columns)?; + let result = filter + .expression() + .evaluate(&filter_batch)? + .into_array(num_outer_rows)?; + let bool_arr = result + .as_any() + .downcast_ref::<BooleanArray>() + .ok_or_else(|| { + DataFusionError::Internal( + "Filter expression did not return BooleanArray".to_string(), + ) + })?; + // Treat nulls as false + if bool_arr.null_count() > 0 { + Ok(arrow::compute::prep_null_mask_filter(bool_arr)) + } else { + Ok(bool_arr.clone()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/exec.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/exec.rs new file mode 100644 index 0000000..b489055 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/exec.rs
@@ -0,0 +1,826 @@ +// 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. + +//! Defines the Sort-Merge join execution plan. +//! A Sort-Merge join plan consumes two sorted children plans and produces +//! joined output by given join type and other options. + +use std::fmt::Formatter; +use std::sync::Arc; + +use super::bitwise_stream::BitwiseSortMergeJoinStream; +use super::materializing_stream::MaterializingSortMergeJoinStream; +use super::metrics::SortMergeJoinMetrics; +use crate::execution_plan::{EmissionType, boundedness_from_children}; +use crate::expressions::PhysicalSortExpr; +use crate::joins::utils::{ + JoinFilter, JoinOn, JoinOnRef, build_join_schema, check_join_is_valid, + estimate_join_statistics, reorder_output_after_swap, + symmetric_join_output_partitioning, +}; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet, SpillMetrics}; +use crate::projection::{ + ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, + physical_to_column_exprs, update_join_on, +}; +use crate::spill::spill_manager::SpillManager; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, validate_child_count, +}; + +use arrow::compute::SortOptions; +use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, internal_err, + plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_physical_expr::equivalence::join_equivalence_properties; +use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; + +/// Join execution plan that executes equi-join predicates on multiple partitions using Sort-Merge +/// join algorithm and applies an optional filter post join. Can be used to join arbitrarily large +/// inputs where one or both of the inputs don't fit in the available memory. +/// +/// # Join Expressions +/// +/// Equi-join predicate (e.g. `<col1> = <col2>`) expressions are represented by [`Self::on`]. +/// +/// Non-equality predicates, which can not be pushed down to join inputs (e.g. +/// `<col1> != <col2>`) are known as "filter expressions" and are evaluated +/// after the equijoin predicates. They are represented by [`Self::filter`]. These are optional +/// expressions. +/// +/// # Sorting +/// +/// Assumes that both the left and right input to the join are pre-sorted. It is not the +/// responsibility of this execution plan to sort the inputs. +/// +/// # "Streamed" vs "Buffered" +/// +/// The number of record batches of streamed input currently present in the memory will depend +/// on the output batch size of the execution plan. There is no spilling support for streamed input. +/// The comparisons are performed from values of join keys in streamed input with the values of +/// join keys in buffered input. One row in streamed record batch could be matched with multiple rows in +/// buffered input batches. Streamed input batches are represented by `StreamedBatch`. +/// +/// Buffered input is buffered for all record batches having the same value of join key. +/// If the memory limit increases beyond the specified value and spilling is enabled, +/// buffered batches could be spilled to disk. If spilling is disabled, the execution +/// will fail under the same conditions. Multiple record batches of buffered could currently reside +/// in memory/disk during the execution. The number of buffered batches residing in +/// memory/disk depends on the number of rows of buffered input having the same value +/// of join key as that of streamed input rows currently present in memory. Due to pre-sorted inputs, +/// the algorithm understands when it is not needed anymore, and releases the buffered batches +/// from memory/disk. Buffered input batches are represented by `BufferedBatch`. +/// +/// Depending on the type of join, left or right input may be selected as streamed or buffered +/// respectively. For example, in a left-outer join, the left execution plan will be selected as +/// streamed input while in a right-outer join, the right execution plan will be selected as the +/// streamed input. +/// +/// Reference for the algorithm: +/// <https://en.wikipedia.org/wiki/Sort-merge_join>. +/// +/// Helpful short video demonstration: +/// <https://www.youtube.com/watch?v=jiWCPJtDE2c>. +#[derive(Debug, Clone)] +pub struct SortMergeJoinExec { + /// Left sorted joining execution plan + pub left: Arc<dyn ExecutionPlan>, + /// Right sorting joining execution plan + pub right: Arc<dyn ExecutionPlan>, + /// Set of common columns used to join on + pub on: JoinOn, + /// Filters which are applied while finding matching rows + pub filter: Option<JoinFilter>, + /// How the join is performed + pub join_type: JoinType, + /// The schema once the join is applied + schema: SchemaRef, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// The left SortExpr + left_sort_exprs: LexOrdering, + /// The right SortExpr + right_sort_exprs: LexOrdering, + /// Sort options of join columns used in sorting left and right execution plans + pub sort_options: Vec<SortOptions>, + /// Defines the null equality for the join. + pub null_equality: NullEquality, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl SortMergeJoinExec { + /// Tries to create a new [SortMergeJoinExec]. + /// The inputs are sorted using `sort_options` are applied to the columns in the `on` + /// # Error + /// This function errors when it is not possible to join the left and right sides on keys `on`. + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: Option<JoinFilter>, + join_type: JoinType, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, + ) -> Result<Self> { + let left_schema = left.schema(); + let right_schema = right.schema(); + + check_join_is_valid(&left_schema, &right_schema, &on)?; + if sort_options.len() != on.len() { + return plan_err!( + "Expected number of sort options: {}, actual: {}", + on.len(), + sort_options.len() + ); + } + + let (left_sort_exprs, right_sort_exprs): (Vec<_>, Vec<_>) = on + .iter() + .zip(sort_options.iter()) + .map(|((l, r), sort_op)| { + let left = PhysicalSortExpr { + expr: Arc::clone(l), + options: *sort_op, + }; + let right = PhysicalSortExpr { + expr: Arc::clone(r), + options: *sort_op, + }; + (left, right) + }) + .unzip(); + let Some(left_sort_exprs) = LexOrdering::new(left_sort_exprs) else { + return plan_err!( + "SortMergeJoinExec requires valid sort expressions for its left side" + ); + }; + let Some(right_sort_exprs) = LexOrdering::new(right_sort_exprs) else { + return plan_err!( + "SortMergeJoinExec requires valid sort expressions for its right side" + ); + }; + + let schema = + Arc::new(build_join_schema(&left_schema, &right_schema, &join_type).0); + let cache = + Self::compute_properties(&left, &right, Arc::clone(&schema), join_type, &on)?; + Ok(Self { + left, + right, + on, + filter, + join_type, + schema, + metrics: ExecutionPlanMetricsSet::new(), + left_sort_exprs, + right_sort_exprs, + sort_options, + null_equality, + cache: Arc::new(cache), + }) + } + + /// Get probe side (e.g streaming side) information for this sort merge join. + /// In current implementation, probe side is determined according to join type. + pub fn probe_side(join_type: &JoinType) -> JoinSide { + // When output schema contains only the right side, probe side is right. + // Otherwise probe side is the left side. + match join_type { + // TODO: sort merge support for right mark (tracked here: https://github.com/apache/datafusion/issues/16226) + JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark => JoinSide::Right, + JoinType::Inner + | JoinType::Left + | JoinType::Full + | JoinType::LeftAnti + | JoinType::LeftSemi + | JoinType::LeftMark => JoinSide::Left, + } + } + + /// Calculate order preservation flags for this sort merge join. + fn maintains_input_order(join_type: JoinType) -> Vec<bool> { + match join_type { + JoinType::Inner => vec![true, false], + JoinType::Left + | JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::LeftMark => vec![true, false], + JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark => { + vec![false, true] + } + _ => vec![false, false], + } + } + + /// Set of common columns used to join on + pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] { + &self.on + } + + /// Ref to right execution plan + pub fn right(&self) -> &Arc<dyn ExecutionPlan> { + &self.right + } + + /// Join type + pub fn join_type(&self) -> JoinType { + self.join_type + } + + /// Ref to left execution plan + pub fn left(&self) -> &Arc<dyn ExecutionPlan> { + &self.left + } + + /// Ref to join filter + pub fn filter(&self) -> &Option<JoinFilter> { + &self.filter + } + + /// Ref to sort options + pub fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } + + /// Null equality + pub fn null_equality(&self) -> NullEquality { + self.null_equality + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + join_type: JoinType, + join_on: JoinOnRef, + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &join_type, + schema, + &Self::maintains_input_order(join_type), + Some(Self::probe_side(&join_type)), + join_on, + )?; + + let output_partitioning = + symmetric_join_output_partitioning(left, right, &join_type)?; + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + boundedness_from_children([left, right]), + )) + } + + /// # Notes: + /// + /// This function should be called BEFORE inserting any repartitioning + /// operators on the join's children. Check [`super::super::HashJoinExec::swap_inputs`] + /// for more details. + pub fn swap_inputs(&self) -> Result<Arc<dyn ExecutionPlan>> { + let left = self.left(); + let right = self.right(); + let new_join = SortMergeJoinExec::try_new( + Arc::clone(right), + Arc::clone(left), + self.on() + .iter() + .map(|(l, r)| (Arc::clone(r), Arc::clone(l))) + .collect::<Vec<_>>(), + self.filter().as_ref().map(JoinFilter::swap), + self.join_type().swap(), + self.sort_options.clone(), + self.null_equality, + )?; + + // TODO: OR this condition with having a built-in projection (like + // ordinary hash join) when we support it. + if matches!( + self.join_type(), + JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ) { + Ok(Arc::new(new_join)) + } else { + reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) + } + } +} + +impl DisplayAs for SortMergeJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let on = self + .on + .iter() + .map(|(c1, c2)| format!("({c1}, {c2})")) + .collect::<Vec<String>>() + .join(", "); + let display_null_equality = + if self.null_equality() == NullEquality::NullEqualsNull { + ", NullsEqual: true" + } else { + "" + }; + write!( + f, + "{}: join_type={:?}, on=[{}]{}{}", + Self::static_name(), + self.join_type, + on, + self.filter.as_ref().map_or_else( + || "".to_string(), + |f| format!(", filter={}", f.expression()) + ), + display_null_equality, + ) + } + DisplayFormatType::TreeRender => { + let on = self + .on + .iter() + .map(|(c1, c2)| { + format!("({} = {})", fmt_sql(c1.as_ref()), fmt_sql(c2.as_ref())) + }) + .collect::<Vec<String>>() + .join(", "); + + if self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + writeln!(f, "on={on}")?; + + if self.null_equality() == NullEquality::NullEqualsNull { + writeln!(f, "NullsEqual: true")?; + } + + Ok(()) + } + } + } +} + +impl ExecutionPlan for SortMergeJoinExec { + fn name(&self) -> &'static str { + "SortMergeJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + let (left_expr, right_expr) = self + .on + .iter() + .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) + .unzip(); + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![ + Some(OrderingRequirements::from(self.left_sort_exprs.clone())), + Some(OrderingRequirements::from(self.right_sort_exprs.clone())), + ] + } + + fn maintains_input_order(&self) -> Vec<bool> { + Self::maintains_input_order(self.join_type) + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + let filter = self.filter.iter().map(|filter| filter.expression()); + crate::apply_expression_roots(join_keys.chain(filter), f) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.filter.clone(), + self.join_type, + self.sort_options.clone(), + self.null_equality, + )?)), + _ => internal_err!("SortMergeJoin wrong number of children"), + }, + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + left_partitions, + right_partitions, + "Invalid SortMergeJoinExec, partition count mismatch {left_partitions}!={right_partitions},\ + consider using RepartitionExec" + ); + let (on_left, on_right) = self.on.iter().cloned().unzip(); + let (streamed, buffered, on_streamed, on_buffered) = + if SortMergeJoinExec::probe_side(&self.join_type) == JoinSide::Left { + ( + Arc::clone(&self.left), + Arc::clone(&self.right), + on_left, + on_right, + ) + } else { + ( + Arc::clone(&self.right), + Arc::clone(&self.left), + on_right, + on_left, + ) + }; + + // execute children plans + let streamed = streamed.execute(partition, Arc::clone(&context))?; + let buffered = buffered.execute(partition, Arc::clone(&context))?; + + let batch_size = context.session_config().batch_size(); + let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]")) + .register(context.memory_pool()); + let spill_manager = SpillManager::new( + context.runtime_env(), + SpillMetrics::new(&self.metrics, partition), + buffered.schema(), + ) + .with_compression_type(context.session_config().spill_compression()); + + if matches!( + self.join_type, + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark + ) { + BitwiseSortMergeJoinStream::try_new( + Arc::clone(&self.schema), + self.sort_options.clone(), + self.null_equality, + streamed, + buffered, + on_streamed, + on_buffered, + self.filter.clone(), + self.join_type, + batch_size, + partition, + &self.metrics, + reservation, + spill_manager, + context.runtime_env(), + ) + } else { + MaterializingSortMergeJoinStream::try_new( + Arc::clone(&self.schema), + self.sort_options.clone(), + self.null_equality, + streamed, + buffered, + on_streamed, + on_buffered, + self.filter.clone(), + self.join_type, + batch_size, + SortMergeJoinMetrics::new(partition, &self.metrics), + reservation, + spill_manager, + context.runtime_env(), + ) + } + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition), ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + // SortMergeJoinExec uses symmetric hash partitioning where both left and right + // inputs are hash-partitioned on the join keys. This means partition `i` of the + // left input is joined with partition `i` of the right input. + // + // TODO stats: it is not possible in general to know the output size of joins + // There are some special cases though, for example: + // - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)` + let left_stats = input_stats[0].as_ref().clone(); + let right_stats = input_stats[1].as_ref().clone(); + Ok(Arc::new(estimate_join_statistics( + left_stats, + right_stats, + &self.on, + self.null_equality, + &self.join_type, + &self.schema, + )?)) + } + + /// Tries to swap the projection with its input [`SortMergeJoinExec`]. If it can be done, + /// it returns the new swapped version having the [`SortMergeJoinExec`] as the top plan. + /// Otherwise, it returns None. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // Convert projected PhysicalExpr's to columns. If not possible, we cannot proceed. + let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) + else { + return Ok(None); + }; + + let (far_right_left_col_ind, far_left_right_col_ind) = join_table_borders( + self.left().schema().fields().len(), + &projection_as_columns, + ); + + if !join_allows_pushdown( + &projection_as_columns, + &self.schema(), + far_right_left_col_ind, + far_left_right_col_ind, + ) { + return Ok(None); + } + + let Some(new_on) = update_join_on( + &projection_as_columns[0..=far_right_left_col_ind as _], + &projection_as_columns[far_left_right_col_ind as _..], + self.on(), + self.left().schema().fields().len(), + ) else { + return Ok(None); + }; + + let (new_left, new_right) = new_join_children( + &projection_as_columns, + far_right_left_col_ind, + far_left_right_col_ind, + self.children()[0], + self.children()[1], + )?; + + Ok(Some(Arc::new(SortMergeJoinExec::try_new( + Arc::new(new_left), + Arc::new(new_right), + new_on, + self.filter.clone(), + self.join_type, + self.sort_options.clone(), + self.null_equality, + )?))) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + let on = self + .on() + .iter() + .map(|(left, right)| { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(left)?), + right: Some(ctx.encode_expr(right)?), + }) + }) + .collect::<Result<Vec<_>>>()?; + + let join_type = crate::joins::proto::join_type_to_proto(self.join_type()); + let null_equality = + crate::joins::proto::null_equality_to_proto(self.null_equality()); + let filter = self + .filter() + .as_ref() + .map(|filter| crate::joins::proto::join_filter_to_proto(filter, ctx)) + .transpose()?; + let sort_options = self + .sort_options() + .iter() + .map(|options| protobuf::SortExprNode { + expr: None, + asc: !options.descending, + nulls_first: options.nulls_first, + }) + .collect(); + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin(Box::new( + protobuf::SortMergeJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + filter, + sort_options, + null_equality: null_equality.into(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortMergeJoinExec { + /// Reconstruct a [`SortMergeJoinExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let sort_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin, + "SortMergeJoinExec", + ); + let left = ctx.decode_required_child( + sort_join.left.as_deref(), + "SortMergeJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sort_join.right.as_deref(), + "SortMergeJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = sort_join + .on + .iter() + .map(|columns| { + let left = ctx.decode_required_expr( + columns.left.as_ref(), + left_schema.as_ref(), + "SortMergeJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + columns.right.as_ref(), + right_schema.as_ref(), + "SortMergeJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::<Result<JoinOn>>()?; + + let join_type = crate::joins::proto::join_type_from_proto( + sort_join.join_type, + "SortMergeJoinExec", + )?; + let null_equality = crate::joins::proto::null_equality_from_proto( + sort_join.null_equality, + "SortMergeJoinExec", + )?; + let filter = sort_join + .filter + .as_ref() + .map(|filter| { + crate::joins::proto::join_filter_from_proto( + filter, + ctx, + "SortMergeJoinExec", + ) + }) + .transpose()?; + let sort_options = sort_join + .sort_options + .iter() + .map(|options| SortOptions { + descending: !options.asc, + nulls_first: options.nulls_first, + }) + .collect(); + + Ok(Arc::new(Self::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )?)) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/filter.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/filter.rs new file mode 100644 index 0000000..4fc6ccc --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/filter.rs
@@ -0,0 +1,388 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Filter handling for Sort-Merge Join +//! +//! This module encapsulates the complexity of join filter evaluation, including: +//! - Immediate filtering for INNER joins +//! - Deferred filtering for outer joins +//! - Metadata tracking for grouping output rows by input row +//! - Correcting filter masks to handle multiple matches per input row + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayBuilder, ArrayRef, BooleanArray, BooleanBuilder, RecordBatch, + RecordBatchOptions, UInt64Array, UInt64Builder, new_null_array, +}; +use arrow::compute::kernels::zip::zip; +use arrow::compute::{self, filter_record_batch}; +use arrow::datatypes::SchemaRef; +use datafusion_common::{JoinSide, JoinType, Result}; + +use crate::joins::utils::JoinFilter; + +/// Metadata for tracking filter results during deferred filtering +/// +/// When a join filter is present and we need to ensure each input row produces +/// at least one output (outer joins), we can't filter immediately. Instead, +/// we accumulate all joined rows with metadata, then post-process to determine +/// which rows to output. +#[derive(Debug)] +pub struct FilterMetadata { + /// Did each output row pass the join filter? + /// Used to detect if an input row found ANY match + pub filter_mask: BooleanBuilder, + + /// Which input row (within batch) produced each output row? + /// Used for grouping output rows by input row + pub row_indices: UInt64Builder, + + /// Which input batch did each output row come from? + /// Used to disambiguate row_indices across multiple batches + pub batch_ids: Vec<usize>, +} + +impl FilterMetadata { + /// Create new empty filter metadata + pub fn new() -> Self { + Self { + filter_mask: BooleanBuilder::new(), + row_indices: UInt64Builder::new(), + batch_ids: vec![], + } + } + + /// Returns (row_indices, filter_mask, batch_ids_ref) and clears builders + pub fn finish_metadata(&mut self) -> (UInt64Array, BooleanArray, &[usize]) { + let row_indices = self.row_indices.finish(); + let filter_mask = self.filter_mask.finish(); + (row_indices, filter_mask, &self.batch_ids) + } + + /// Add metadata for null-joined rows (no filter applied) + pub fn append_nulls(&mut self, num_rows: usize) { + self.filter_mask.append_nulls(num_rows); + self.row_indices.append_nulls(num_rows); + self.batch_ids.resize( + self.batch_ids.len() + num_rows, + 0, // batch_id = 0 for null-joined rows + ); + } + + /// Add metadata for filtered rows + pub fn append_filter_metadata( + &mut self, + row_indices: &UInt64Array, + filter_mask: &BooleanArray, + batch_id: usize, + ) { + debug_assert_eq!( + row_indices.len(), + filter_mask.len(), + "row_indices and filter_mask must have same length" + ); + + self.filter_mask.extend(filter_mask); + self.row_indices.extend(row_indices); + self.batch_ids + .resize(self.batch_ids.len() + row_indices.len(), batch_id); + } + + /// Verify that metadata arrays are aligned (same length) + pub fn debug_assert_metadata_aligned(&self) { + if self.filter_mask.len() > 0 { + debug_assert_eq!( + self.filter_mask.len(), + self.row_indices.len(), + "filter_mask and row_indices must have same length when metadata is used" + ); + debug_assert_eq!( + self.filter_mask.len(), + self.batch_ids.len(), + "filter_mask and batch_ids must have same length when metadata is used" + ); + } else { + debug_assert_eq!( + self.filter_mask.len(), + 0, + "filter_mask should be empty when batches is empty" + ); + } + } +} + +impl Default for FilterMetadata { + fn default() -> Self { + Self::new() + } +} + +/// Determines if a join type needs deferred filtering +/// +/// Deferred filtering is required when: +/// - A filter exists AND +/// - The join type requires ensuring each input row produces at least one output +pub fn needs_deferred_filtering( + filter: &Option<JoinFilter>, + join_type: JoinType, +) -> bool { + filter.is_some() + && matches!(join_type, JoinType::Left | JoinType::Right | JoinType::Full) +} + +/// Gets the arrays which join filters are applied on +/// +/// Extracts the columns needed for filter evaluation from left and right batch columns +pub fn get_filter_columns( + join_filter: &Option<JoinFilter>, + left_columns: &[ArrayRef], + right_columns: &[ArrayRef], +) -> Vec<ArrayRef> { + let mut filter_columns = vec![]; + + if let Some(f) = join_filter { + let left_columns: Vec<ArrayRef> = f + .column_indices() + .iter() + .filter(|col_index| col_index.side == JoinSide::Left) + .map(|i| Arc::clone(&left_columns[i.index])) + .collect(); + let right_columns: Vec<ArrayRef> = f + .column_indices() + .iter() + .filter(|col_index| col_index.side == JoinSide::Right) + .map(|i| Arc::clone(&right_columns[i.index])) + .collect(); + + filter_columns.extend(left_columns); + filter_columns.extend(right_columns); + } + + filter_columns +} + +/// Determines if current index is the last occurrence of a row +/// +/// Used during filter mask correction to detect row boundaries when grouping +/// output rows by input row. +fn last_index_for_row( + row_index: usize, + indices: &UInt64Array, + batch_ids: &[usize], + indices_len: usize, +) -> bool { + debug_assert_eq!( + indices.len(), + indices_len, + "indices.len() should match indices_len parameter" + ); + debug_assert_eq!( + batch_ids.len(), + indices_len, + "batch_ids.len() should match indices_len" + ); + debug_assert!( + row_index < indices_len, + "row_index {row_index} should be < indices_len {indices_len}", + ); + + // If this is the last index overall, it's definitely the last for this row + if row_index == indices_len - 1 { + return true; + } + + // Check if next row has different (batch_id, index) pair + let current_batch_id = batch_ids[row_index]; + let next_batch_id = batch_ids[row_index + 1]; + + if current_batch_id != next_batch_id { + return true; + } + + // Same batch_id, check if row index is different + // Both current and next should be non-null (already joined rows) + if indices.is_null(row_index) || indices.is_null(row_index + 1) { + return true; + } + + indices.value(row_index) != indices.value(row_index + 1) +} + +/// Corrects the filter mask for joins with deferred filtering +/// +/// When an input row joins with multiple buffered rows, we get multiple output rows. +/// This function groups them by input row and applies join-type-specific logic: +/// +/// - **Outer joins**: Keep first matching row, convert rest to nulls, add null-joined for unmatched +/// +/// # Arguments +/// * `join_type` - The type of join being performed +/// * `row_indices` - Which input row produced each output row +/// * `batch_ids` - Which batch each output row came from +/// * `filter_mask` - Whether each output row passed the filter +/// * `expected_size` - Total number of input rows (for adding unmatched) +/// +/// # Returns +/// Corrected mask indicating which rows to include in final output: +/// - `true`: Include this row +/// - `false`: Convert to null-joined row (outer joins) +/// - `null`: Discard this row +pub fn get_corrected_filter_mask( + join_type: JoinType, + row_indices: &UInt64Array, + batch_ids: &[usize], + filter_mask: &BooleanArray, + expected_size: usize, +) -> Option<BooleanArray> { + let row_indices_length = row_indices.len(); + let mut corrected_mask: BooleanBuilder = + BooleanBuilder::with_capacity(row_indices_length); + let mut seen_true = false; + + match join_type { + JoinType::Left | JoinType::Right | JoinType::Full => { + // For each input row group: keep first filter-passing row, + // discard (null) remaining matches, null-join if none passed. + // Null metadata entries are already-null-joined rows that + // flow through unchanged to preserve output ordering. + for i in 0..row_indices_length { + let last_index = + last_index_for_row(i, row_indices, batch_ids, row_indices_length); + if filter_mask.is_null(i) { + corrected_mask.append_value(true); + } else if filter_mask.value(i) { + seen_true = true; + corrected_mask.append_value(true); + } else if seen_true || !filter_mask.value(i) && !last_index { + corrected_mask.append_null(); + } else { + corrected_mask.append_value(false); + } + + if last_index { + seen_true = false; + } + } + + corrected_mask.append_n(expected_size - corrected_mask.len(), false); + Some(corrected_mask.finish()) + } + JoinType::LeftMark + | JoinType::RightMark + | JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftAnti + | JoinType::RightAnti => { + unreachable!("Semi/anti/mark joins are handled by BitwiseSortMergeJoinStream") + } + JoinType::Inner => None, + } +} + +/// Applies corrected filter mask to record batch based on join type +/// +/// The corrected mask has three possible values per row: +/// - `true`: Keep the row as-is (matched and passed filter) +/// - `false`: Convert to null-joined row (all filter matches failed for this input row) +/// - `null`: Discard the row entirely (duplicate match for an already-output input row) +/// +/// This function preserves input row ordering by processing each row in place +/// rather than separating matched/unmatched rows. +pub fn filter_record_batch_by_join_type( + record_batch: &RecordBatch, + corrected_mask: &BooleanArray, + join_type: JoinType, + schema: &SchemaRef, + buffered_schema: &SchemaRef, +) -> Result<RecordBatch> { + match join_type { + JoinType::Left | JoinType::Right | JoinType::Full => { + if record_batch.num_rows() == 0 { + return Ok(record_batch.clone()); + } + + // Discard null-masked rows (keep true + false only) + let keep_mask = compute::is_not_null(corrected_mask)?; + let kept_batch = filter_record_batch(record_batch, &keep_mask)?; + + if kept_batch.num_rows() == 0 { + return Ok(kept_batch); + } + + let kept_corrected = compute::filter(corrected_mask, &keep_mask)?; + let kept_corrected = kept_corrected + .as_any() + .downcast_ref::<BooleanArray>() + .unwrap(); + + // All rows passed the filter — no null-joining needed + if !kept_corrected.has_false() { + return Ok(kept_batch); + } + + // For false entries: replace the non-preserved side with nulls. + // This preserves row ordering unlike filter+concat. + let (null_side_start, null_side_len) = match join_type { + JoinType::Left => { + // Left join: null out right (buffered) columns + let left_cols = + schema.fields().len() - buffered_schema.fields().len(); + (left_cols, buffered_schema.fields().len()) + } + JoinType::Right => { + // Right join: null out left (buffered) columns + (0, buffered_schema.fields().len()) + } + JoinType::Full => { + // Full join: null out buffered columns for streamed rows + // that matched but failed the filter. Unmatched buffered + // rows are null-joined on the streamed side separately + // when the buffered batch is drained. + let left_cols = + schema.fields().len() - buffered_schema.fields().len(); + (left_cols, buffered_schema.fields().len()) + } + _ => unreachable!(), + }; + + let num_rows = kept_batch.num_rows(); + let mut columns: Vec<ArrayRef> = kept_batch.columns().to_vec(); + + for col in columns.iter_mut().skip(null_side_start).take(null_side_len) { + let null_array = new_null_array(col.data_type(), num_rows); + *col = zip(kept_corrected, &*col, &null_array)?; + } + + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + Ok(RecordBatch::try_new_with_options( + Arc::clone(schema), + columns, + &options, + )?) + } + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::LeftMark + | JoinType::RightMark => unreachable!( + "Semi/anti/mark joins are handled by SemiAntiMarkSortMergeJoinStream" + ), + JoinType::Inner => Ok(filter_record_batch(record_batch, corrected_mask)?), + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/materializing_stream.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/materializing_stream.rs new file mode 100644 index 0000000..3baa0c4 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/materializing_stream.rs
@@ -0,0 +1,2000 @@ +// 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 execution +//! +//! This module implements the Sort-Merge Join operator as an async +//! generator running a merge scan: it drives two sorted input streams (the +//! *streamed* side and the *buffered* side), compares join keys, and +//! produces joined `RecordBatch`es. + +use std::cmp::Ordering; +use std::collections::{HashMap, VecDeque}; +use std::fmt::Debug; +use std::mem::size_of; +use std::ops::Range; +use std::sync::Arc; + +use crate::joins::sort_merge_join::filter::{ + FilterMetadata, filter_record_batch_by_join_type, get_corrected_filter_mask, + get_filter_columns, needs_deferred_filtering, +}; +use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; +use crate::joins::utils::{JoinFilter, JoinKeyComparator}; +use crate::metrics::Time; +use crate::spill::spill_manager::SpillManager; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; +use crate::{PhysicalExpr, SendableRecordBatchStream}; + +use arrow::array::{types::UInt64Type, *}; +use arrow::compute::{ + self, BatchCoalescer, SortOptions, concat_batches, filter_record_batch, interleave, + take_arrays, +}; +use arrow::datatypes::SchemaRef; +use datafusion_common::cast::as_uint64_array; +use datafusion_common::instant::Instant; +use datafusion_common::{ + DataFusionError, JoinType, NullEquality, Result, exec_err, internal_err, +}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_execution::{SpillFile, TryEmitter, async_try_stream}; +use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; + +use futures::StreamExt; + +/// Represents a chunk of joined data from streamed and buffered side +pub(super) struct StreamedJoinedChunk { + /// Index of batch in buffered_data + buffered_batch_idx: Option<usize>, + /// Array builder for streamed indices + streamed_indices: UInt64Builder, + /// Array builder for buffered indices + /// This could contain nulls if the join is null-joined + buffered_indices: UInt64Builder, +} + +/// Represents a record batch from streamed input. +/// +/// Also stores information of matching rows from buffered batches. +pub(super) struct StreamedBatch { + /// The streamed record batch + pub batch: RecordBatch, + /// The index of row in the streamed batch to compare with buffered batches + pub idx: usize, + /// The join key arrays of streamed batch which are used to compare with buffered batches + /// and to produce output. They are produced by evaluating `on` expressions. + pub join_arrays: Vec<ArrayRef>, + /// Chunks of indices from buffered side (may be nulls) joined to streamed + pub output_indices: Vec<StreamedJoinedChunk>, + /// Total number of output rows across all chunks in `output_indices` + pub num_output_rows: usize, + /// Index of currently scanned batch from buffered data + pub buffered_batch_idx: Option<usize>, +} + +impl StreamedBatch { + fn new(batch: RecordBatch, on_column: &[Arc<dyn PhysicalExpr>]) -> Self { + let join_arrays = join_arrays(&batch, on_column); + StreamedBatch { + batch, + idx: 0, + join_arrays, + output_indices: vec![], + num_output_rows: 0, + buffered_batch_idx: None, + } + } + + fn new_empty(schema: SchemaRef) -> Self { + StreamedBatch { + batch: RecordBatch::new_empty(schema), + idx: 0, + join_arrays: vec![], + output_indices: vec![], + num_output_rows: 0, + buffered_batch_idx: None, + } + } + + /// Number of unfrozen output pairs in this streamed batch + fn num_output_rows(&self) -> usize { + self.num_output_rows + } + + /// Appends new pair consisting of current streamed index and `buffered_idx` + /// index of buffered batch with `buffered_batch_idx` index. + fn append_output_pair( + &mut self, + buffered_batch_idx: Option<usize>, + buffered_idx: Option<usize>, + batch_size: usize, + ) { + // If no current chunk exists or current chunk is not for current buffered batch, + // create a new chunk + if self.output_indices.is_empty() || self.buffered_batch_idx != buffered_batch_idx + { + // Compute capacity only when creating a new chunk (infrequent operation). + // The capacity is the remaining space to reach batch_size. + // This should always be >= 1 since we only call this when num_output_rows < batch_size. + debug_assert!( + batch_size > self.num_output_rows, + "batch_size ({batch_size}) must be > num_output_rows ({})", + self.num_output_rows + ); + let capacity = batch_size - self.num_output_rows; + self.output_indices.push(StreamedJoinedChunk { + buffered_batch_idx, + streamed_indices: UInt64Builder::with_capacity(capacity), + buffered_indices: UInt64Builder::with_capacity(capacity), + }); + self.buffered_batch_idx = buffered_batch_idx; + }; + let current_chunk = self.output_indices.last_mut().unwrap(); + + // Append index of streamed batch and index of buffered batch into current chunk + current_chunk.streamed_indices.append_value(self.idx as u64); + if let Some(idx) = buffered_idx { + current_chunk.buffered_indices.append_value(idx as u64); + } else { + current_chunk.buffered_indices.append_null(); + } + self.num_output_rows += 1; + } +} + +/// Per-row filter outcome tracking for full outer joins. +/// +/// In a full outer join with a filter, buffered rows that match on join +/// keys but fail every filter evaluation must be emitted with NULLs on +/// the streamed side. Three states are needed because a simple boolean +/// cannot distinguish "never matched" (handled by [`BufferedBatch::null_joined`]) +/// from "matched but all filters failed" (must be emitted as null-joined). +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum FilterState { + /// Row never appeared in a matched pair. + Unvisited = 0, + /// Row matched streamed rows, but all filter evaluations failed. + AllFailed = 1, + /// Row matched and at least one filter evaluation passed. + SomePassed = 2, +} + +/// A buffered batch that contains contiguous rows with same join key +/// +/// `BufferedBatch` can exist as either an in-memory `RecordBatch` or a `SpillFile`. +#[derive(Debug)] +pub(super) struct BufferedBatch { + /// Represents in memory or spilled record batch + pub batch: BufferedBatchState, + /// The range in which the rows share the same join key + pub range: Range<usize>, + /// Array refs of the join key + pub join_arrays: Vec<ArrayRef>, + /// Buffered joined index (null joining buffered) + pub null_joined: Vec<usize>, + /// Size estimation used for reserving / releasing memory + pub size_estimation: usize, + /// Memory footprint of `join_arrays` cached at construction time. + /// Used during spill to track the residual memory that remains after + /// the main batch is written to disk. + pub join_arrays_mem: usize, + /// Actual amount tracked in the memory reservation for this batch. + /// + /// - `InMemory`: equals `size_estimation` (full batch + join_arrays + metadata) + /// - `Spilled`: equals `join_arrays_mem` (join key arrays stay in memory) + /// + /// Invariant: `free_reservation()` shrinks by exactly this amount, so we never + /// shrink by more than we grew. + pub reserved_amount: usize, + /// Tracks filter outcomes for buffered rows in full outer joins. + /// Indexed by absolute row position within the batch. See [`FilterState`]. + pub join_filter_status: Vec<FilterState>, + /// Current buffered batch number of rows. Equal to batch.num_rows() + /// but if batch is spilled to disk this property is preferable + /// and less expensive + pub num_rows: usize, +} + +impl BufferedBatch { + fn new( + batch: RecordBatch, + range: Range<usize>, + on_column: &[PhysicalExprRef], + ) -> Self { + let join_arrays = join_arrays(&batch, on_column); + + // Estimation is calculated as + // inner batch size + // + join keys size + // + worst case null_joined (as vector capacity * element size) + // + Range size + // + size of this estimation + let join_arrays_mem: usize = join_arrays + .iter() + .map(|arr| arr.get_array_memory_size()) + .sum(); + + let size_estimation = batch.get_array_memory_size() + + join_arrays_mem + + batch.num_rows().next_power_of_two() * size_of::<usize>() + + size_of::<Range<usize>>() + + size_of::<usize>(); + + let num_rows = batch.num_rows(); + BufferedBatch { + batch: BufferedBatchState::InMemory(batch), + range, + join_arrays, + null_joined: vec![], + size_estimation, + join_arrays_mem, + reserved_amount: 0, + join_filter_status: vec![FilterState::Unvisited; num_rows], + num_rows, + } + } +} + +// TODO: Spill join arrays (https://github.com/apache/datafusion/pull/17429) +// Used to represent whether the buffered data is currently in memory or written to disk +pub(super) enum BufferedBatchState { + // In memory record batch + InMemory(RecordBatch), + // Spilled temp file + Spilled(Arc<dyn SpillFile>), +} + +impl Debug for BufferedBatchState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InMemory(batch) => f.debug_tuple("InMemory").field(batch).finish(), + Self::Spilled(_) => { + write!(f, "Spilled(Custom_Backend)") + } + } + } +} +/// Sort-Merge join stream for Inner/Left/Right/Full joins. +/// +/// Named "materializing" because it builds explicit `(streamed, buffered)` row +/// pairs in [`JoinedRecordBatches`] to produce output columns from both sides +/// of the join. +pub(super) struct MaterializingSortMergeJoinStream { + // ======================================================================== + // PROPERTIES: + // These fields are initialized at the start and remain constant throughout + // the execution. + // ======================================================================== + /// Output schema + pub schema: SchemaRef, + /// Defines the null equality for the join. + pub null_equality: NullEquality, + /// Sort options of join columns used to sort streamed and buffered data stream + pub sort_options: Vec<SortOptions>, + /// optional join filter + pub filter: Option<JoinFilter>, + /// How the join is performed + pub join_type: JoinType, + /// Cached `needs_deferred_filtering(filter, join_type)` — both inputs + /// are fixed at construction time. + pub deferred_filtering: bool, + /// Target output batch size + pub batch_size: usize, + + // ======================================================================== + // STREAMED FIELDS: + // These fields manage the properties and state of the streamed input. + // ======================================================================== + /// Input schema of streamed + pub streamed_schema: SchemaRef, + /// Streamed data stream + pub streamed: SendableRecordBatchStream, + /// Current processing record batch of streamed + pub streamed_batch: StreamedBatch, + /// True once the streamed input has no more rows + pub streamed_exhausted: bool, + /// Join key columns of streamed + pub on_streamed: Vec<PhysicalExprRef>, + + // ======================================================================== + // BUFFERED FIELDS: + // These fields manage the properties and state of the buffered input. + // ======================================================================== + /// Input schema of buffered + pub buffered_schema: SchemaRef, + /// Buffered data stream + pub buffered: SendableRecordBatchStream, + /// Current buffered data + pub buffered_data: BufferedData, + /// Has any streamed row matched the current buffered key group? + /// (FULL join: an unmatched group is emitted null-joined when passed.) + pub buffered_group_matched: bool, + /// True once the buffered input has no more rows and no group remains + pub buffered_exhausted: bool, + /// Join key columns of buffered + pub on_buffered: Vec<PhysicalExprRef>, + + // ======================================================================== + // MERGE JOIN STATES: + // These fields track the execution state of merge join and are updated + // during the execution. + // ======================================================================== + /// Staging output array builders + pub joined_record_batches: JoinedRecordBatches, + /// Output buffer. Currently used by filtering as it requires double buffering + /// to avoid small/empty batches. Non-filtered joins output directly from + /// `joined_record_batches.joined_batches` + pub output: BatchCoalescer, + /// Manages the process of spilling and reading back intermediate data + pub spill_manager: SpillManager, + + /// Tracks the number of batches currently spilled + pub spilled_batch_count: usize, + + /// Time spent doing the join's own work (including spill write and + /// read-back). The clock is stopped while awaiting the child inputs or + /// the consumer taking an emitted batch — see [`Self::stop_join_time`]. + pub join_time: Time, + /// Start of the currently running `join_time` span; `None` while the + /// clock is stopped. + pub join_time_start: Option<Instant>, + + // ======================================================================== + // CACHED COMPARATORS: + // Pre-built comparators to avoid per-row type dispatch in hot loops. + // ======================================================================== + /// Comparator for streamed vs buffered head batch key comparison + pub streamed_buffered_cmp: Option<JoinKeyComparator>, + /// Comparator for buffered head vs tail batch equality check + pub buffered_equality_cmp: Option<JoinKeyComparator>, + + // ======================================================================== + // EXECUTION RESOURCES: + // Fields related to managing execution resources and monitoring performance. + // ======================================================================== + /// Metrics + pub join_metrics: SortMergeJoinMetrics, + /// Memory reservation + pub reservation: MemoryReservation, + /// Runtime env + pub runtime_env: Arc<RuntimeEnv>, + /// A unique id per streamed batch, tagging deferred-filter metadata so + /// `get_corrected_filter_mask` can group output rows by input batch. + pub streamed_batch_counter: usize, +} + +/// Staging area for joined data before output +/// +/// Accumulates joined rows until either: +/// - Target batch size reached (for efficiency) +/// - Stream exhausted (flush remaining data) +pub(super) struct JoinedRecordBatches { + /// Joined batches. Each batch is already joined columns from left and right sources + pub(super) joined_batches: BatchCoalescer, + /// Filter metadata for deferred filtering + pub(super) filter_metadata: FilterMetadata, +} + +impl JoinedRecordBatches { + /// Concatenates all accumulated batches into a single RecordBatch + /// + /// Must drain ALL batches from BatchCoalescer for filtered joins to ensure + /// metadata alignment when applying get_corrected_filter_mask(). + pub(super) fn concat_batches(&mut self, schema: &SchemaRef) -> Result<RecordBatch> { + self.joined_batches.finish_buffered_batch()?; + + let mut all_batches = vec![]; + while let Some(batch) = self.joined_batches.next_completed_batch() { + all_batches.push(batch); + } + + match all_batches.as_slice() { + [] => unreachable!("concat_batches called with empty BatchCoalescer"), + [single_batch] => Ok(single_batch.clone()), + multiple_batches => Ok(concat_batches(schema, multiple_batches)?), + } + } + + /// Clears batches without touching metadata (for early return when no filtering needed) + fn clear_batches(&mut self, schema: &SchemaRef, batch_size: usize) { + self.joined_batches = BatchCoalescer::new(Arc::clone(schema), batch_size) + .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)); + } + + /// Asserts that if batches is empty, metadata is also empty + #[inline] + fn debug_assert_empty_consistency(&self) { + if self.joined_batches.is_empty() { + debug_assert_eq!( + self.filter_metadata.filter_mask.len(), + 0, + "filter_mask should be empty when batches is empty" + ); + debug_assert_eq!( + self.filter_metadata.row_indices.len(), + 0, + "row_indices should be empty when batches is empty" + ); + debug_assert_eq!( + self.filter_metadata.batch_ids.len(), + 0, + "batch_ids should be empty when batches is empty" + ); + } + } + + /// Pushes a batch with null metadata (rows that need no filter correction) + /// + /// Used for: (1) Full join buffered rows with no streamed match, and + /// (2) outer join streamed rows with no buffered match. These rows are + /// already in final form but must flow through the deferred filtering + /// pipeline to preserve output ordering. Null metadata causes + /// get_corrected_filter_mask() to pass them through unchanged. + /// + /// Maintains invariant: N rows → N metadata entries (nulls) + fn push_batch_with_null_metadata(&mut self, batch: RecordBatch, join_type: JoinType) { + debug_assert!( + matches!(join_type, JoinType::Left | JoinType::Right | JoinType::Full), + "push_batch_with_null_metadata should only be called for deferred-filtered joins" + ); + + let num_rows = batch.num_rows(); + + self.filter_metadata.append_nulls(num_rows); + + self.filter_metadata.debug_assert_metadata_aligned(); + self.joined_batches + .push_batch(batch) + .expect("Failed to push batch to BatchCoalescer"); + } + + /// Pushes a batch with filter metadata (filtered outer joins) + /// + /// Deferred filtering: An input row may join with multiple buffered rows, but we + /// don't know yet if all matches failed the filter. We track metadata so + /// `get_corrected_filter_mask()` can later group by input row and decide: + /// - If any match passed: emit passing rows + /// - If all matches failed: emit null-joined row + /// + /// Maintains invariant: N rows → N metadata entries + fn push_batch_with_filter_metadata( + &mut self, + batch: RecordBatch, + row_indices: &UInt64Array, + filter_mask: &BooleanArray, + streamed_batch_id: usize, + join_type: JoinType, + ) { + debug_assert!( + matches!(join_type, JoinType::Left | JoinType::Right | JoinType::Full), + "push_batch_with_filter_metadata should only be called for outer joins that need deferred filtering" + ); + + debug_assert_eq!( + row_indices.len(), + filter_mask.len(), + "row_indices and filter_mask must have same length" + ); + + self.filter_metadata.append_filter_metadata( + row_indices, + filter_mask, + streamed_batch_id, + ); + + self.filter_metadata.debug_assert_metadata_aligned(); + self.joined_batches + .push_batch(batch) + .expect("Failed to push batch to BatchCoalescer"); + } + + /// Pushes a batch without metadata (non-filtered joins) + /// + /// No deferred filtering needed. Either every join match is output (Inner), + /// or null-joined rows are handled separately. No need to track which input + /// row produced which output row. + fn push_batch_without_metadata(&mut self, batch: RecordBatch) { + self.joined_batches + .push_batch(batch) + .expect("Failed to push batch to BatchCoalescer"); + } + + fn clear(&mut self, schema: &SchemaRef, batch_size: usize) { + self.joined_batches = BatchCoalescer::new(Arc::clone(schema), batch_size) + .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)); + self.filter_metadata = FilterMetadata::new(); + self.debug_assert_empty_consistency(); + } +} + +impl MaterializingSortMergeJoinStream { + #[expect(clippy::too_many_arguments)] + pub fn try_new( + schema: SchemaRef, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, + streamed: SendableRecordBatchStream, + buffered: SendableRecordBatchStream, + on_streamed: Vec<Arc<dyn PhysicalExpr>>, + on_buffered: Vec<Arc<dyn PhysicalExpr>>, + filter: Option<JoinFilter>, + join_type: JoinType, + batch_size: usize, + join_metrics: SortMergeJoinMetrics, + reservation: MemoryReservation, + spill_manager: SpillManager, + runtime_env: Arc<RuntimeEnv>, + ) -> Result<SendableRecordBatchStream> { + let streamed_schema = streamed.schema(); + let buffered_schema = buffered.schema(); + debug_assert!( + matches!( + join_type, + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full + ), + "MaterializingSortMergeJoinStream does not handle {join_type:?}; \ + semi/anti/mark joins use BitwiseSortMergeJoinStream" + ); + let join_time = join_metrics.join_time(); + let mut this = Self { + sort_options, + null_equality, + schema: Arc::clone(&schema), + streamed_schema: Arc::clone(&streamed_schema), + buffered_schema, + streamed, + buffered, + streamed_batch: StreamedBatch::new_empty(streamed_schema), + buffered_data: BufferedData::default(), + buffered_group_matched: false, + streamed_exhausted: false, + buffered_exhausted: false, + on_streamed, + on_buffered, + deferred_filtering: needs_deferred_filtering(&filter, join_type), + filter, + joined_record_batches: JoinedRecordBatches { + joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size) + .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), + filter_metadata: FilterMetadata::new(), + }, + output: BatchCoalescer::new(schema, batch_size) + .with_biggest_coalesce_batch_size(Option::from(batch_size / 2)), + batch_size, + join_type, + join_metrics, + reservation, + runtime_env, + spill_manager, + spilled_batch_count: 0, + join_time, + join_time_start: None, + streamed_buffered_cmp: None, + buffered_equality_cmp: None, + streamed_batch_counter: 0, + }; + + let schema = Arc::clone(&this.schema); + let baseline_metrics = this.join_metrics.baseline_metrics(); + + let stream = async_try_stream(|mut emitter| async move { + this.start_join_time(); + let result = this.join(&mut emitter).await; + this.stop_join_time(); + result + }); + // ObservedStream records the baseline metrics (output rows/batches, + // end time). + Ok(Box::pin(ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, + None, + ))) + } + + /// Main loop: the textbook sort-merge join. + /// + /// Both inputs arrive sorted on the join keys. The streamed side is + /// consumed one row at a time; the buffered side one key *group* (all + /// contiguous rows sharing a key) at a time + async fn join( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + // 1. Load the first streamed row and the first buffered key group. + self.load_next_streamed_batch().await?; + self.advance_buffered_group().await?; + + // 2. Merge-scan while either input still has rows. + while !(self.streamed_exhausted && self.buffered_exhausted) { + // Flush the deferred-filtering pipeline once a full batch of + // rows accumulated (filtered outer joins output through it). + if self.deferred_filtering + && self.deferred_rows_accumulated() >= self.batch_size + { + self.emit_deferred_output(emitter).await?; + } + + // 3. Compare the join keys at both cursors. An exhausted side + // compares as the larger one, so the other side keeps + // draining through its own arm. + match self.compare_streamed_buffered()? { + // 3a. The streamed row can never match: null-join it (outer + // joins emit it; inner joins drop it), then advance. + Ordering::Less => { + self.null_join_streamed_row(); + if self.num_unfrozen_pairs() >= self.batch_size { + self.freeze_and_emit(emitter).await?; + } + if !self.try_advance_streamed_row() { + self.load_next_streamed_batch().await?; + } + } + // 3b. The buffered group can never match again: null-join + // it if nothing matched it (FULL join), then advance to + // the next key group. + Ordering::Greater => { + self.null_join_buffered_group(); + if !self.try_advance_buffered_group()? { + self.advance_buffered_group().await?; + } + } + // 3c. Match: pair the streamed row with the whole group — + // materializing ("freezing") mid-scan whenever a full + // batch of pairs accumulates — then advance streamed. + // The group stays for the next streamed row. + Ordering::Equal => { + while !self.pair_streamed_row_with_group() { + self.freeze_and_emit(emitter).await?; + } + if !self.try_advance_streamed_row() { + self.load_next_streamed_batch().await?; + } + } + } + + // 4. Emit completed output batches (filtered joins emit + // through the deferred-filtering pipeline above instead). + if !self.deferred_filtering + && self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + self.emit_completed_joined_batches(emitter).await; + } + } + + // 5. Flush everything that remains. + self.on_children_exhausted(emitter).await + } + + /// `Equal`: pair the current streamed row with every row of the + /// buffered key group, and mark the group as matched. + /// + /// Returns false when a full batch of pairs has accumulated (the scan + /// may or may not be complete): the caller must materialize + /// (`freeze_and_emit`) and call again, which resumes the scan where it + /// paused. Returns true when the group scan is complete and there is + /// room for more pairs. + fn pair_streamed_row_with_group(&mut self) -> bool { + while !self.buffered_data.scanning_finished() + && self.num_unfrozen_pairs() < self.batch_size + { + let scanning_idx = self.buffered_data.scanning_idx(); + self.streamed_batch.append_output_pair( + Some(self.buffered_data.scanning_batch_idx), + Some(scanning_idx), + self.batch_size, + ); + self.buffered_data.scanning_advance(); + } + if self.num_unfrozen_pairs() >= self.batch_size { + return false; + } + + self.buffered_group_matched = true; + self.buffered_data.scanning_reset(); + true + } + + /// `Less` (outer joins): no buffered row matches the current streamed + /// row — emit it joined to NULLs. Inner joins emit nothing. + fn null_join_streamed_row(&mut self) { + if matches!( + self.join_type, + JoinType::Left | JoinType::Right | JoinType::Full + ) { + let scanning_batch_idx = if self.buffered_data.scanning_finished() { + None + } else { + Some(self.buffered_data.scanning_batch_idx) + }; + self.streamed_batch.append_output_pair( + scanning_batch_idx, + None, + self.batch_size, + ); + } + self.buffered_data.scanning_reset(); + } + + /// `Greater` (FULL join): the buffered group can never match a streamed + /// row anymore — if nothing matched it, mark all its rows for + /// null-joined output (produced when the group's batches are dequeued). + fn null_join_buffered_group(&mut self) { + if self.join_type == JoinType::Full && !self.buffered_group_matched { + while !self.buffered_data.scanning_finished() { + let scanning_idx = self.buffered_data.scanning_idx(); + self.buffered_data + .scanning_batch_mut() + .null_joined + .push(scanning_idx); + self.buffered_data.scanning_advance(); + } + } + self.buffered_data.scanning_reset(); + } + + /// Start (resume) the `join_time` clock. + fn start_join_time(&mut self) { + debug_assert!(self.join_time_start.is_none(), "join_time already running"); + self.join_time_start = Some(Instant::now()); + } + + /// Stop (pause) the `join_time` clock, accumulating the elapsed span. + /// + /// Called around awaits whose duration is not the join's own work: the + /// child input streams' `next()` and `emitter.emit()` (where the + /// consumer processes the batch). The join's own spill write and + /// read-back are NOT excluded — that time is join work. + fn stop_join_time(&mut self) { + if let Some(start) = self.join_time_start.take() { + self.join_time.add_elapsed(start); + } + } + + /// Number of rows currently waiting in the deferred-filtering pipeline. + /// + /// Typically bounded to ~2*batch_size: one batch_size worth from + /// freeze_dequeuing_buffered() (when an input batch is fully consumed), + /// plus up to batch_size pairs accumulating toward the next freeze. A + /// single streamed row matching a very large key group can exceed that + /// (its pairs freeze into the pipeline before the gate runs again — same + /// as the pre-generator design). This does not reintroduce the unbounded + /// buffering fixed by PR #20482; `on_children_exhausted` flushes the + /// remainder. + fn deferred_rows_accumulated(&self) -> usize { + self.num_unfrozen_pairs() + + self.joined_record_batches.filter_metadata.filter_mask.len() + } + + /// Run the deferred-filtering pipeline over everything accumulated so + /// far and emit its completed output, if any. Clears the accumulation + /// it processed. + /// + /// The caller gates this on `deferred_rows_accumulated() >= batch_size`: + /// running the pipeline per row instead (concat + correct_mask + + /// filter_by_type) would dominate runtime for unique keys. + async fn emit_deferred_output( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + // Ensure required spilled batches are restored to memory before + // processing, as this path invokes freeze_all(). + self.restore_spilled_batches_for_freeze().await?; + if let Some(batch) = self.process_filtered_batches()? { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); + emitter.emit(batch).await; + self.start_join_time(); + } + Ok(()) + } + + /// Restore every spilled buffered batch that the next freeze needs. + async fn restore_spilled_batches_for_freeze(&mut self) -> Result<()> { + let needed = self.get_required_batch_indices(self.buffered_data.batches.len()); + self.restore_spilled_batches(&needed).await + } + + /// Emit all completed joined batches to the stream consumer. + async fn emit_completed_joined_batches( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) { + while let Some(record_batch) = self + .joined_record_batches + .joined_batches + .next_completed_batch() + { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); + emitter.emit(record_batch).await; + self.start_join_time(); + } + } + + /// Flush everything that remains once both inputs are exhausted. + async fn on_children_exhausted( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + // Freeze the remaining pairs, restoring any spilled batches needed. + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; + + // Verify metadata alignment before final output + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + if self.deferred_filtering { + // Filtered joins must concat and filter ALL remaining data at once + if !self.joined_record_batches.joined_batches.is_empty() { + let record_batch = self.filter_joined_batch()?; + self.stop_join_time(); + emitter.emit(record_batch).await; + self.start_join_time(); + } + } else if !self.joined_record_batches.joined_batches.is_empty() { + // For non-filtered joins, finish buffered data first, then emit + // every completed batch. + self.joined_record_batches + .joined_batches + .finish_buffered_batch()?; + self.emit_completed_joined_batches(emitter).await; + } + + // Drain the double-buffering coalescer used by filtered joins. + if !self.output.is_empty() { + self.output.finish_buffered_batch()?; + while let Some(record_batch) = self.output.next_completed_batch() { + self.stop_join_time(); + emitter.emit(record_batch).await; + self.start_join_time(); + } + } + + Ok(()) + } + + /// Build a comparator for streamed vs buffered head batch keys. + fn rebuild_streamed_buffered_cmp(&mut self) -> Result<()> { + if self.streamed_batch.join_arrays.is_empty() + || !self.buffered_data.has_buffered_rows() + { + self.streamed_buffered_cmp = None; + return Ok(()); + } + self.streamed_buffered_cmp = Some(JoinKeyComparator::new( + &self.streamed_batch.join_arrays, + &self.buffered_data.head_batch().join_arrays, + &self.sort_options, + self.null_equality, + )?); + Ok(()) + } + + /// Build a comparator for buffered head vs tail batch equality. + fn rebuild_buffered_equality_cmp(&mut self) -> Result<()> { + if self.buffered_data.batches.is_empty() { + self.buffered_equality_cmp = None; + return Ok(()); + } + self.buffered_equality_cmp = Some(JoinKeyComparator::new( + &self.buffered_data.head_batch().join_arrays, + &self.buffered_data.tail_batch().join_arrays, + &self.sort_options, + // is_join_arrays_equal treats both-null as equal + NullEquality::NullEqualsNull, + )?); + Ok(()) + } + + /// Number of unfrozen output pairs (used to decide when to freeze + output) + fn num_unfrozen_pairs(&self) -> usize { + self.streamed_batch.num_output_rows() + } + + /// Process accumulated batches for filtered joins + /// + /// Freezes unfrozen pairs, applies deferred filtering, and returns a + /// completed output batch if one is ready. + fn process_filtered_batches(&mut self) -> Result<Option<RecordBatch>> { + self.freeze_all()?; + + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + if !self.joined_record_batches.joined_batches.is_empty() { + let out_filtered_batch = self.filter_joined_batch()?; + self.output + .push_batch(out_filtered_batch) + .expect("Failed to push output batch"); + + if self.output.has_completed_batch() { + let record_batch = self + .output + .next_completed_batch() + .expect("Failed to get output batch"); + return Ok(Some(record_batch)); + } + } + + Ok(None) + } + + /// Identifies which buffered batches are needed for the upcoming freeze operation + fn get_required_batch_indices(&self, buffered_freeze_count: usize) -> Vec<usize> { + let mut needed = vec![]; + // Avoid scanning if no spilled batches exist + if self.spilled_batch_count == 0 { + return needed; + } + // We need all batches that matched with streamed rows + for chunk in &self.streamed_batch.output_indices { + if let Some(idx) = chunk.buffered_batch_idx { + needed.push(idx); + } + } + + // Full Joins need to emit null-joined rows, so we need batches up to freeze_count + if self.join_type == JoinType::Full { + needed.extend(0..buffered_freeze_count); + } + + needed.sort_unstable(); + needed.dedup(); + needed + } + + /// Asynchronously reads spilled batches back into memory. + /// Only processes the required indices to avoid OOMs. + async fn restore_spilled_batches( + &mut self, + required_indices: &[usize], + ) -> Result<()> { + for &idx in required_indices { + // Guard against indices that might be out of bounds if the queue was cleared + if idx >= self.buffered_data.batches.len() { + continue; + } + + let bb = &mut self.buffered_data.batches[idx]; + + if let BufferedBatchState::Spilled(spill_file) = &bb.batch { + let mut spill_stream = self + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; + + match spill_stream.next().await.transpose()? { + Some(batch) => { + // Transition the batch back to InMemory + bb.batch = BufferedBatchState::InMemory(batch); + self.spilled_batch_count -= 1; + // The batch is back in memory, so we must account for its size. + let newly_allocated = + bb.size_estimation.saturating_sub(bb.reserved_amount); + self.reservation.grow(newly_allocated); + bb.reserved_amount = bb.size_estimation; + + self.join_metrics + .peak_mem_used() + .set_max(self.reservation.size()); + } + None => { + return internal_err!("Spill file was empty"); + } + } + } + } + + Ok(()) + } + + /// Sync fast path of advancing the streamed cursor: move to the next row + /// of the current batch. Returns false at the batch boundary, where the + /// caller must load the next batch via + /// [`Self::load_next_streamed_batch`]. + fn try_advance_streamed_row(&mut self) -> bool { + if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() { + self.streamed_batch.idx += 1; + return true; + } + false + } + + /// Load the next streamed batch (freezing the finished one) and point + /// the streamed cursor at its first row. Sets `streamed_exhausted` when + /// the streamed input has no more rows. + async fn load_next_streamed_batch(&mut self) -> Result<()> { + loop { + // Loading a new streamed batch freezes the current one, which + // materializes buffered columns — restore any spilled buffered + // batches it needs first. + self.restore_spilled_batches_for_freeze().await?; + + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.streamed.next().await.transpose(); + self.start_join_time(); + match item? { + None => { + // Release the streamed input pipeline's resources. + let streamed_schema = self.streamed.schema(); + self.streamed = + Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + self.streamed_exhausted = true; + return Ok(()); + } + Some(batch) => { + if batch.num_rows() > 0 { + self.freeze_streamed()?; + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + self.streamed_batch = + StreamedBatch::new(batch, &self.on_streamed); + self.rebuild_streamed_buffered_cmp()?; + // Every incoming streamed batch gets a unique id. + self.streamed_batch_counter += 1; + return Ok(()); + } + } + } + } + } + + fn free_reservation(&mut self, buffered_batch: &BufferedBatch) { + if buffered_batch.reserved_amount > 0 { + self.reservation.shrink(buffered_batch.reserved_amount); + } + } + + fn allocate_reservation(&mut self, mut buffered_batch: BufferedBatch) -> Result<()> { + match self.reservation.try_grow(buffered_batch.size_estimation) { + Ok(_) => { + buffered_batch.reserved_amount = buffered_batch.size_estimation; + self.join_metrics + .peak_mem_used() + .set_max(self.reservation.size()); + Ok(()) + } + Err(_) if self.runtime_env.disk_manager.tmp_files_enabled() => { + // Spill buffered batch to disk + + match buffered_batch.batch { + BufferedBatchState::InMemory(batch) => { + let spill_file = self + .spill_manager + .spill_record_batch_and_finish( + &[batch], + "sort_merge_join_buffered_spill", + )? + .unwrap(); // Operation only return None if no batches are spilled, here we ensure that at least one batch is spilled + + buffered_batch.batch = BufferedBatchState::Spilled(spill_file); + self.spilled_batch_count += 1; + + // Join key arrays remain in memory after the batch is + // spilled — the comparator needs them for key boundary + // detection. Force-grow the reservation so the pool + // reflects actual memory usage even if this pushes + // pool.reserved() above the configured limit. This is + // safe because the memory is physically consumed and + // not tracking it would let other operators over-allocate + // against a stale pool view. + let join_arrays_mem = buffered_batch.join_arrays_mem; + self.reservation.grow(join_arrays_mem); + buffered_batch.reserved_amount = join_arrays_mem; + self.join_metrics + .peak_mem_used() + .set_max(self.reservation.size()); + + Ok(()) + } + _ => internal_err!("Buffered batch has empty body"), + } + } + Err(e) => exec_err!("{}. Disk spilling disabled.", e.message()), + }?; + + self.buffered_data.batches.push_back(buffered_batch); + Ok(()) + } + + /// Sync fast path of [`Self::advance_buffered_group`]: when the next + /// group starts in the single remaining buffered batch and provably ends + /// within it (the common case — a group only reaches a batch boundary + /// once per batch), advance entirely synchronously. Returns false — + /// leaving all state unchanged — when the async path must run instead. + fn try_advance_buffered_group(&mut self) -> Result<bool> { + if self.buffered_data.batches.len() != 1 { + return Ok(false); + } + let head_batch = self.buffered_data.head_batch(); + if head_batch.range.end == head_batch.num_rows { + // Fully consumed — needs dequeuing (and loading the next batch). + return Ok(false); + } + + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; + } + let cmp = self.buffered_equality_cmp.as_ref().unwrap(); + + // Scan the next group's extent before committing any state, so a + // bail-out (the group may span into the next batch) leaves + // everything untouched for the async path. + let batch = self.buffered_data.head_batch(); + let group_start = batch.range.end; + let mut group_end = group_start + 1; + while group_end < batch.num_rows && cmp.is_equal(group_start, group_end) { + group_end += 1; + } + if group_end == batch.num_rows { + return Ok(false); + } + + let batch = self.buffered_data.tail_batch_mut(); + batch.range.start = group_start; + batch.range.end = group_end; + self.buffered_group_matched = false; + Ok(true) + } + + /// Advance the buffered side to the next key group: dequeue batches + /// fully consumed by the previous group, then collect all contiguous + /// rows sharing the next join key (the group may span multiple buffered + /// batches). Sets `buffered_exhausted` when no group remains. + async fn advance_buffered_group(&mut self) -> Result<()> { + self.buffered_group_matched = false; + self.dequeue_consumed_buffered_batches().await?; + + if self.buffered_data.batches.is_empty() { + // Load the batch holding the first row of the next group. + if !self.load_next_buffered_batch().await? { + self.buffered_exhausted = true; + return Ok(()); + } + } else { + // Seed the next group at the first unconsumed row of the + // remaining batch. + let tail_batch = self.buffered_data.tail_batch_mut(); + tail_batch.range.start = tail_batch.range.end; + tail_batch.range.end += 1; + } + + self.extend_buffered_group().await + } + + /// Dequeue buffered batches fully consumed by the previous group, + /// producing their pending output (e.g. Full-join null-joined rows). + async fn dequeue_consumed_buffered_batches(&mut self) -> Result<()> { + let mut head_changed = false; + while !self.buffered_data.batches.is_empty() { + let head_batch = self.buffered_data.head_batch(); + if head_batch.range.end != head_batch.num_rows { + // The next group starts within the head batch: streamed rows + // will be joined with the head batch in the next step. + break; + } + // load the spilled head batch before dequeuing + let needed = self.get_required_batch_indices(1); + self.restore_spilled_batches(&needed).await?; + + self.freeze_dequeuing_buffered()?; + if let Some(mut buffered_batch) = self.buffered_data.batches.pop_front() { + self.produce_buffered_not_matched(&mut buffered_batch)?; + self.free_reservation(&buffered_batch); + if matches!(buffered_batch.batch, BufferedBatchState::Spilled(_)) { + self.spilled_batch_count -= 1; + } + head_changed = true; + } + } + if head_changed { + self.streamed_buffered_cmp = None; + self.buffered_equality_cmp = None; + } + Ok(()) + } + + /// Load the next non-empty buffered batch and seed a new group with its + /// first row. Returns false when the buffered input is exhausted. + async fn load_next_buffered_batch(&mut self) -> Result<bool> { + loop { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.buffered.next().await.transpose(); + self.start_join_time(); + match item? { + None => { + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + return Ok(false); + } + Some(batch) => { + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..1, &self.on_buffered); + self.allocate_reservation(buffered_batch)?; + self.streamed_buffered_cmp = None; + return Ok(true); + } + } + } + } + } + + /// Extend the current group with every following row that shares its + /// key, loading more buffered batches as needed. + async fn extend_buffered_group(&mut self) -> Result<()> { + loop { + if self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; + } + while self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.as_ref().unwrap().is_equal( + self.buffered_data.head_batch().range.start, + self.buffered_data.tail_batch().range.end, + ) { + self.buffered_data.tail_batch_mut().range.end += 1; + } else { + // Group complete within the current batch. + return Ok(()); + } + } + } else { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.buffered.next().await.transpose(); + self.start_join_time(); + match item? { + None => { + // Group complete; the input is done but the group is + // still valid — `buffered_exhausted` is only set once + // it has been fully consumed and dequeued. + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + return Ok(()); + } + Some(batch) => { + // Polling batches coming concurrently as multiple partitions + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..0, &self.on_buffered); + self.allocate_reservation(buffered_batch)?; + self.buffered_equality_cmp = None; + } + } + } + } + } + } + + /// Get comparison result of streamed row and buffered batches + fn compare_streamed_buffered(&mut self) -> Result<Ordering> { + if self.streamed_exhausted { + return Ok(Ordering::Greater); + } + if !self.buffered_data.has_buffered_rows() { + return Ok(Ordering::Less); + } + + if self.streamed_buffered_cmp.is_none() { + self.rebuild_streamed_buffered_cmp()?; + } + Ok(self.streamed_buffered_cmp.as_ref().unwrap().compare( + self.streamed_batch.idx, + self.buffered_data.head_batch().range.start, + )) + } + + /// Materialize ("freeze") the accumulated pairs — restoring any spilled + /// batches they reference first — and emit completed output batches + /// (filtered joins emit through the deferred-filtering gate instead). + async fn freeze_and_emit( + &mut self, + emitter: &mut TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; + + if !self.deferred_filtering + && self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + self.emit_completed_joined_batches(emitter).await; + } + Ok(()) + } + + fn freeze_all(&mut self) -> Result<()> { + self.freeze_buffered(self.buffered_data.batches.len())?; + self.freeze_streamed()?; + + // After freezing, metadata should be aligned + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + Ok(()) + } + + // Produces and stages record batches to ensure dequeued buffered batch + // no longer needed: + // 1. freezes all indices joined to streamed side + // 2. freezes NULLs joined to dequeued buffered batch to "release" it + fn freeze_dequeuing_buffered(&mut self) -> Result<()> { + self.freeze_streamed()?; + // Only freeze and produce the first batch in buffered_data as the batch is fully processed + self.freeze_buffered(1)?; + + // After freezing, metadata should be aligned + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + Ok(()) + } + + // Produces and stages record batch from buffered indices with corresponding + // NULLs on streamed side. + // + // Applicable only in case of Full join. + // + fn freeze_buffered(&mut self, batch_count: usize) -> Result<()> { + if self.join_type != JoinType::Full { + return Ok(()); + } + for buffered_batch in self.buffered_data.batches.range_mut(..batch_count) { + let buffered_indices = UInt64Array::from_iter_values( + buffered_batch.null_joined.iter().map(|&index| index as u64), + ); + if let Some(record_batch) = produce_buffered_null_batch( + &self.schema, + &self.streamed_schema, + &buffered_indices, + buffered_batch, + )? { + self.joined_record_batches + .push_batch_with_null_metadata(record_batch, self.join_type); + } + buffered_batch.null_joined.clear(); + } + Ok(()) + } + + fn produce_buffered_not_matched( + &mut self, + buffered_batch: &mut BufferedBatch, + ) -> Result<()> { + if self.join_type != JoinType::Full { + return Ok(()); + } + + // Collect buffered rows that matched on join keys but had every + // filter evaluation fail — these must be emitted with NULLs on + // the streamed side to satisfy full outer join semantics. + let not_matched_buffered_indices = buffered_batch + .join_filter_status + .iter() + .enumerate() + .filter_map(|(i, state)| { + matches!(state, FilterState::AllFailed).then_some(i as u64) + }) + .collect::<Vec<_>>(); + + let buffered_indices = + UInt64Array::from_iter_values(not_matched_buffered_indices.iter().copied()); + + if let Some(record_batch) = produce_buffered_null_batch( + &self.schema, + &self.streamed_schema, + &buffered_indices, + buffered_batch, + )? { + self.joined_record_batches + .push_batch_with_null_metadata(record_batch, self.join_type); + } + buffered_batch + .join_filter_status + .fill(FilterState::Unvisited); + + Ok(()) + } + + // Produces and stages record batch for all output indices found + // for current streamed batch and clears staged output indices. + // + // Null-joined chunks (no buffered match) are pushed immediately. + // Matched chunks are collected and processed together in + // freeze_streamed_matched() to amortize filter evaluation overhead. + fn freeze_streamed(&mut self) -> Result<()> { + let mut matched_chunks: Vec<(usize, UInt64Array, UInt64Array)> = Vec::new(); + let mut total_matched_rows: usize = 0; + + for chunk in self.streamed_batch.output_indices.iter_mut() { + let left_indices = chunk.streamed_indices.finish(); + if left_indices.is_empty() { + continue; + } + let right_indices: UInt64Array = chunk.buffered_indices.finish(); + + if chunk.buffered_batch_idx.is_none() { + let left_columns = + materialize_left_columns(&self.streamed_batch.batch, &left_indices)?; + let right_columns = + create_unmatched_columns(&self.buffered_schema, left_indices.len()); + + let columns = if self.join_type != JoinType::Right { + [left_columns, right_columns].concat() + } else { + [right_columns, left_columns].concat() + }; + let batch = RecordBatch::try_new(Arc::clone(&self.schema), columns)?; + + // Null-joined rows (no buffered match) need no filter correction, + // but must flow through the same pipeline as matched rows to + // preserve output ordering. Use null metadata as a sentinel so + // get_corrected_filter_mask() passes them through unchanged. + if self.deferred_filtering { + self.joined_record_batches + .push_batch_with_null_metadata(batch, self.join_type); + } else { + self.joined_record_batches + .push_batch_without_metadata(batch); + } + continue; + } + + total_matched_rows += left_indices.len(); + matched_chunks.push(( + chunk.buffered_batch_idx.unwrap(), + left_indices, + right_indices, + )); + } + + if !matched_chunks.is_empty() { + self.freeze_streamed_matched(&matched_chunks, total_matched_rows)?; + } + + self.streamed_batch.output_indices.clear(); + self.streamed_batch.num_output_rows = 0; + Ok(()) + } + + /// Materializes columns, evaluates the join filter, and pushes output + /// for all matched chunks in a single batch. This avoids per-chunk + /// RecordBatch construction and filter evaluation, which dominates + /// cost when keys are near-unique (1 row per chunk). + fn freeze_streamed_matched( + &mut self, + matched_chunks: &[(usize, UInt64Array, UInt64Array)], + total_matched_rows: usize, + ) -> Result<()> { + debug_assert!( + !matched_chunks.is_empty(), + "caller guards this with an is_empty check before calling" + ); + debug_assert!( + matched_chunks.iter().all(|(idx, left, right)| { + left.len() == right.len() && *idx < self.buffered_data.batches.len() + }), + "left/right indices are built in pairs from the same streamed×buffered cross, \ + and batch_idx comes from iterating buffered_data.batches" + ); + debug_assert_eq!( + matched_chunks + .iter() + .map(|(_, l, _)| l.len()) + .sum::<usize>(), + total_matched_rows, + "total_matched_rows is accumulated from the same chunks in freeze_streamed" + ); + + let combined_left_indices = if matched_chunks.len() == 1 { + matched_chunks[0].1.clone() + } else { + let refs: Vec<&dyn Array> = + matched_chunks.iter().map(|c| &c.1 as &dyn Array).collect(); + as_uint64_array(&compute::concat(&refs)?)?.clone() + }; + + let left_columns = + materialize_left_columns(&self.streamed_batch.batch, &combined_left_indices)?; + + let right_columns = + self.materialize_right_columns(matched_chunks, total_matched_rows)?; + + let filter_columns = if self.join_type == JoinType::Right { + get_filter_columns(&self.filter, &right_columns, &left_columns) + } else { + get_filter_columns(&self.filter, &left_columns, &right_columns) + }; + + let columns = if self.join_type != JoinType::Right { + [left_columns, right_columns].concat() + } else { + [right_columns, left_columns].concat() + }; + let output_batch = RecordBatch::try_new(Arc::clone(&self.schema), columns)?; + + if !filter_columns.is_empty() { + if let Some(f) = &self.filter { + let filter_batch = + RecordBatch::try_new(Arc::clone(f.schema()), filter_columns)?; + let filter_result = f + .expression() + .evaluate(&filter_batch)? + .into_array(filter_batch.num_rows())?; + + let filter_result_mask = + datafusion_common::cast::as_boolean_array(&filter_result)?; + + // Convert NULL filter results to false — NULL means "not satisfied" + // per SQL semantics, same as Left/Right outer joins. + let mask = if filter_result_mask.null_count() > 0 { + compute::prep_null_mask_filter(filter_result_mask) + } else { + filter_result_mask.clone() + }; + + if self.deferred_filtering { + self.joined_record_batches.push_batch_with_filter_metadata( + output_batch, + &combined_left_indices, + &mask, + self.streamed_batch_counter, + self.join_type, + ); + } else { + let filtered_batch = filter_record_batch(&output_batch, &mask)?; + self.joined_record_batches + .push_batch_without_metadata(filtered_batch); + } + + // Track which buffered rows had all filter matches fail, + // so full join can emit them as null-joined later. + if self.join_type == JoinType::Full { + let mut offset = 0usize; + for (batch_idx, _left, right) in matched_chunks { + let chunk_len = right.len(); + let buffered_batch = &mut self.buffered_data.batches[*batch_idx]; + + for i in 0..chunk_len { + if right.is_null(i) { + continue; + } + let idx = right.value(i) as usize; + match buffered_batch.join_filter_status[idx] { + FilterState::SomePassed => {} + _ if mask.value(offset + i) => { + buffered_batch.join_filter_status[idx] = + FilterState::SomePassed; + } + _ => { + buffered_batch.join_filter_status[idx] = + FilterState::AllFailed; + } + } + } + offset += chunk_len; + } + debug_assert_eq!( + offset, total_matched_rows, + "offset must advance through every chunk exactly once" + ); + } + } + } else { + self.joined_record_batches + .push_batch_without_metadata(output_batch); + } + + Ok(()) + } + + /// Materializes right-side columns across all matched chunks. + /// + /// When chunks reference a single buffered batch, indices are concatenated + /// for a single fetch. When multiple batches are involved, `interleave` + /// gathers columns across sources. A null-row sentinel at source index 0 + /// handles null right indices (unmatched streamed rows). + fn materialize_right_columns( + &mut self, + matched_chunks: &[(usize, UInt64Array, UInt64Array)], + total_matched_rows: usize, + ) -> Result<Vec<ArrayRef>> { + let first_batch_idx = matched_chunks[0].0; + let single_source = matched_chunks.iter().all(|c| c.0 == first_batch_idx); + + if single_source { + let combined_right_indices = if matched_chunks.len() == 1 { + matched_chunks[0].2.clone() + } else { + let refs: Vec<&dyn Array> = + matched_chunks.iter().map(|c| &c.2 as &dyn Array).collect(); + as_uint64_array(&compute::concat(&refs)?)?.clone() + }; + + return fetch_right_columns_by_idxs( + &self.buffered_data, + first_batch_idx, + &combined_right_indices, + ); + } + + // Multiple source batches: map each buffered_batch_idx to a + // contiguous source index, reserving source 0 for a null sentinel. + let mut batch_idx_to_source: HashMap<usize, usize> = HashMap::new(); + let mut source_batches: Vec<usize> = Vec::new(); + for (batch_idx, _, _) in matched_chunks { + batch_idx_to_source.entry(*batch_idx).or_insert_with(|| { + let idx = source_batches.len() + 1; + source_batches.push(*batch_idx); + idx + }); + } + + let mut interleave_indices: Vec<(usize, usize)> = + Vec::with_capacity(total_matched_rows); + for (batch_idx, _, right) in matched_chunks { + let source = batch_idx_to_source[batch_idx]; + for i in 0..right.len() { + if right.is_null(i) { + interleave_indices.push((0, 0)); + } else { + interleave_indices.push((source, right.value(i) as usize)); + } + } + } + + let num_right_cols = self.buffered_schema.fields().len(); + + // Read each source batch once (spilled batches require disk I/O). + let source_data_result: Result<Vec<RecordBatch>> = source_batches + .iter() + .map(|&idx| { + let bb = &self.buffered_data.batches[idx]; + match &bb.batch { + BufferedBatchState::InMemory(batch) => Ok(batch.clone()), + BufferedBatchState::Spilled(_) => { + internal_err!("Buffered batch should have been unspilled before fetching columns") + } + } + }) + .collect(); + + let source_data = source_data_result?; + + let mut right_columns = Vec::with_capacity(num_right_cols); + for col_idx in 0..num_right_cols { + let dtype = self.buffered_schema.field(col_idx).data_type(); + let null_array = new_null_array(dtype, 1); + + let mut source_arrays: Vec<&dyn Array> = + Vec::with_capacity(source_batches.len() + 1); + source_arrays.push(null_array.as_ref()); + + for data in &source_data { + source_arrays.push(data.column(col_idx).as_ref()); + } + right_columns.push(interleave(&source_arrays, &interleave_indices)?); + } + + Ok(right_columns) + } + + fn filter_joined_batch(&mut self) -> Result<RecordBatch> { + // Metadata should be aligned before processing + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + let record_batch = self.joined_record_batches.concat_batches(&self.schema)?; + let (mut out_indices, mut out_mask, mut batch_ids) = + self.joined_record_batches.filter_metadata.finish_metadata(); + let default_batch_ids = vec![0; record_batch.num_rows()]; + + // If only nulls come in and indices sizes doesn't match with expected record batch count + // generate missing indices + // Happens for null joined batches for Full Join + if out_indices.null_count() == out_indices.len() + && out_indices.len() != record_batch.num_rows() + { + out_mask = BooleanArray::from(vec![None; record_batch.num_rows()]); + out_indices = UInt64Array::from(vec![None; record_batch.num_rows()]); + batch_ids = &default_batch_ids; + } + + // After potential reconstruction, metadata should align with batch row count + debug_assert_eq!( + out_indices.len(), + record_batch.num_rows(), + "out_indices length should match record_batch row count" + ); + debug_assert_eq!( + out_mask.len(), + record_batch.num_rows(), + "out_mask length should match record_batch row count (unless empty)" + ); + debug_assert_eq!( + batch_ids.len(), + record_batch.num_rows(), + "batch_ids length should match record_batch row count" + ); + + if out_mask.is_empty() { + self.joined_record_batches + .clear_batches(&self.schema, self.batch_size); + return Ok(record_batch); + } + + // Validate inputs to get_corrected_filter_mask + debug_assert_eq!( + out_indices.len(), + out_mask.len(), + "out_indices and out_mask must have same length for get_corrected_filter_mask" + ); + debug_assert_eq!( + batch_ids.len(), + out_mask.len(), + "batch_ids and out_mask must have same length for get_corrected_filter_mask" + ); + + let maybe_corrected_mask = get_corrected_filter_mask( + self.join_type, + &out_indices, + batch_ids, + &out_mask, + record_batch.num_rows(), + ); + + let corrected_mask = if let Some(ref filtered_join_mask) = maybe_corrected_mask { + filtered_join_mask + } else { + &out_mask + }; + + self.filter_record_batch_by_join_type(&record_batch, corrected_mask) + } + + fn filter_record_batch_by_join_type( + &mut self, + record_batch: &RecordBatch, + corrected_mask: &BooleanArray, + ) -> Result<RecordBatch> { + let filtered_record_batch = filter_record_batch_by_join_type( + record_batch, + corrected_mask, + self.join_type, + &self.schema, + &self.buffered_schema, + )?; + + self.joined_record_batches + .clear(&self.schema, self.batch_size); + + Ok(filtered_record_batch) + } +} + +/// Materialize left (streamed) columns using slice or take. +fn materialize_left_columns( + batch: &RecordBatch, + indices: &UInt64Array, +) -> Result<Vec<ArrayRef>> { + if let Some(range) = is_contiguous_range(indices) { + Ok(batch.slice(range.start, range.len()).columns().to_vec()) + } else { + Ok(take_arrays(batch.columns(), indices, None)?) + } +} + +fn create_unmatched_columns(schema: &SchemaRef, size: usize) -> Vec<ArrayRef> { + schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), size)) + .collect::<Vec<_>>() +} + +fn produce_buffered_null_batch( + schema: &SchemaRef, + streamed_schema: &SchemaRef, + buffered_indices: &PrimitiveArray<UInt64Type>, + buffered_batch: &BufferedBatch, +) -> Result<Option<RecordBatch>> { + if buffered_indices.is_empty() { + return Ok(None); + } + + // Take buffered (right) columns + let right_columns = + fetch_right_columns_from_batch_by_idxs(buffered_batch, buffered_indices)?; + + // Create null streamed (left) columns + let mut left_columns = streamed_schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), buffered_indices.len())) + .collect::<Vec<_>>(); + + left_columns.extend(right_columns); + + Ok(Some(RecordBatch::try_new( + Arc::clone(schema), + left_columns, + )?)) +} + +/// Checks if a `UInt64Array` contains a contiguous ascending range (e.g. \[3,4,5,6\]). +/// Returns `Some(start..start+len)` if so, `None` otherwise. +/// This allows replacing an O(n) `take` with an O(1) `slice`. +#[inline] +fn is_contiguous_range(indices: &UInt64Array) -> Option<Range<usize>> { + if indices.is_empty() || indices.null_count() > 0 { + return None; + } + let values = indices.values(); + let start = values[0]; + let len = values.len() as u64; + // Quick rejection: if last element doesn't match expected, not contiguous + if values[values.len() - 1] != start + len - 1 { + return None; + } + // Verify every element is sequential (handles duplicates and gaps) + for i in 1..values.len() { + if values[i] != start + i as u64 { + return None; + } + } + Some(start as usize..(start + len) as usize) +} + +/// Get `buffered_indices` rows for `buffered_data[buffered_batch_idx]` by specific column indices +#[inline(always)] +fn fetch_right_columns_by_idxs( + buffered_data: &BufferedData, + buffered_batch_idx: usize, + buffered_indices: &UInt64Array, +) -> Result<Vec<ArrayRef>> { + fetch_right_columns_from_batch_by_idxs( + &buffered_data.batches[buffered_batch_idx], + buffered_indices, + ) +} + +#[inline(always)] +fn fetch_right_columns_from_batch_by_idxs( + buffered_batch: &BufferedBatch, + buffered_indices: &UInt64Array, +) -> Result<Vec<ArrayRef>> { + match &buffered_batch.batch { + BufferedBatchState::InMemory(batch) => { + if let Some(range) = is_contiguous_range(buffered_indices) { + Ok(batch.slice(range.start, range.len()).columns().to_vec()) + } else { + Ok(take_arrays(batch.columns(), buffered_indices, None)?) + } + } + BufferedBatchState::Spilled(_) => { + internal_err!( + "Buffered batch should have been unspilled before fetching columns" + ) + } + } +} + +/// Buffered data contains all buffered batches with one unique join key +#[derive(Debug, Default)] +pub(super) struct BufferedData { + /// Buffered batches with the same key + pub batches: VecDeque<BufferedBatch>, + /// current scanning batch index used by the group-scan phase + pub scanning_batch_idx: usize, + /// current scanning offset used by the group-scan phase + pub scanning_offset: usize, +} + +impl BufferedData { + pub fn head_batch(&self) -> &BufferedBatch { + self.batches.front().unwrap() + } + + pub fn tail_batch(&self) -> &BufferedBatch { + self.batches.back().unwrap() + } + + pub fn tail_batch_mut(&mut self) -> &mut BufferedBatch { + self.batches.back_mut().unwrap() + } + + pub fn has_buffered_rows(&self) -> bool { + self.batches.iter().any(|batch| !batch.range.is_empty()) + } + + pub fn scanning_reset(&mut self) { + self.scanning_batch_idx = 0; + self.scanning_offset = 0; + } + + pub fn scanning_advance(&mut self) { + self.scanning_offset += 1; + while !self.scanning_finished() && self.scanning_batch_finished() { + self.scanning_batch_idx += 1; + self.scanning_offset = 0; + } + } + + pub fn scanning_batch(&self) -> &BufferedBatch { + &self.batches[self.scanning_batch_idx] + } + + pub fn scanning_batch_mut(&mut self) -> &mut BufferedBatch { + &mut self.batches[self.scanning_batch_idx] + } + + pub fn scanning_idx(&self) -> usize { + self.scanning_batch().range.start + self.scanning_offset + } + + pub fn scanning_batch_finished(&self) -> bool { + self.scanning_offset == self.scanning_batch().range.len() + } + + pub fn scanning_finished(&self) -> bool { + self.scanning_batch_idx == self.batches.len() + } +} + +/// Get join array refs of given batch and join columns +fn join_arrays(batch: &RecordBatch, on_column: &[PhysicalExprRef]) -> Vec<ArrayRef> { + on_column + .iter() + .map(|c| { + let num_rows = batch.num_rows(); + let c = c.evaluate(batch).unwrap(); + c.into_array(num_rows).unwrap() + }) + .collect() +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/metrics.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/metrics.rs new file mode 100644 index 0000000..6f52a22 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/metrics.rs
@@ -0,0 +1,82 @@ +// 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. + +//! Module for tracking Sort Merge Join metrics + +use crate::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, + MetricCategory, Time, +}; + +/// Metrics for SortMergeJoinExec +pub(super) struct SortMergeJoinMetrics { + /// Total time for joining probe-side batches to the build-side batches + join_time: Time, + /// Number of batches consumed by this operator + input_batches: Count, + /// Number of rows consumed by this operator + input_rows: Count, + /// Execution metrics + baseline_metrics: BaselineMetrics, + /// Peak memory used for buffered data. + /// Calculated as sum of peak memory values across partitions + peak_mem_used: Gauge, +} + +impl SortMergeJoinMetrics { + pub fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition); + let input_batches = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("input_batches", partition); + let input_rows = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("input_rows", partition); + let peak_mem_used = + MetricBuilder::new(metrics).peak_memory_usage("peak_mem_used", partition); + + let baseline_metrics = BaselineMetrics::new(metrics, partition); + + Self { + join_time, + input_batches, + input_rows, + baseline_metrics, + peak_mem_used, + } + } + + pub fn join_time(&self) -> Time { + self.join_time.clone() + } + + pub fn baseline_metrics(&self) -> BaselineMetrics { + self.baseline_metrics.clone() + } + + pub fn input_batches(&self) -> Count { + self.input_batches.clone() + } + + pub fn input_rows(&self) -> Count { + self.input_rows.clone() + } + + pub fn peak_mem_used(&self) -> Gauge { + self.peak_mem_used.clone() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/mod.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/mod.rs new file mode 100644 index 0000000..2fdb092 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/mod.rs
@@ -0,0 +1,29 @@ +// 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 Execution Plan Operator + +pub use exec::SortMergeJoinExec; + +pub(crate) mod bitwise_stream; +mod exec; +mod filter; +pub(crate) mod materializing_stream; +mod metrics; + +#[cfg(test)] +mod tests;
diff --git a/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/tests.rs b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/tests.rs new file mode 100644 index 0000000..175a9c0 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/sort_merge_join/tests.rs
@@ -0,0 +1,5783 @@ +// 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. + +//! SortMergeJoin Testing Module +//! +//! This module currently contains the following test types in this order: +//! - Join behaviour (left, right, full, inner, semi, anti, mark) +//! - Batch spilling +//! - Filter mask +//! +//! Add relevant tests under the specified sections. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::Duration; + +use super::bitwise_stream::BitwiseSortMergeJoinStream; +use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; +use crate::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; +use crate::spill::spill_manager::SpillManager; +use crate::test::TestMemoryExec; +use crate::test::exec::BarrierExec; +use crate::test::{build_table_i32, build_table_i32_two_cols}; +use crate::{ExecutionPlan, RecordBatchStream, common}; +use crate::{ + expressions::Column, joins::sort_merge_join::filter::get_corrected_filter_mask, + joins::sort_merge_join::materializing_stream::JoinedRecordBatches, +}; +use arrow::array::{ + BinaryArray, BooleanArray, Date32Array, Date64Array, FixedSizeBinaryArray, + Int32Array, RecordBatch, UInt64Array, +}; +use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow_ord::sort::SortColumn; +use arrow_schema::SchemaRef; +use bytes::Bytes; +use datafusion_common::JoinType::*; +use datafusion_common::instant::Instant; +use datafusion_common::{ + JoinSide, internal_err, + test_util::{batches_to_sort_string, batches_to_string}, +}; +use datafusion_common::{ + JoinType, NullEquality, Result, ScalarValue, assert_batches_eq, assert_contains, +}; +use datafusion_common_runtime::JoinSet; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, +}; +use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_expr::Operator; +use datafusion_physical_expr::expressions::BinaryExpr; +use datafusion_physical_expr::expressions::Literal; +use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; +use futures::{Stream, StreamExt}; +use insta::assert_snapshot; +use itertools::Itertools; +use std::collections::VecDeque; + +fn build_table( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), +) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +fn build_table_from_batches(batches: Vec<RecordBatch>) -> Arc<dyn ExecutionPlan> { + let schema = batches.first().unwrap().schema(); + TestMemoryExec::try_new_exec(&[batches], schema, None).unwrap() +} + +fn build_date_table( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), +) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Date32, false), + Field::new(b.0, DataType::Date32, false), + Field::new(c.0, DataType::Date32, false), + ]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Date32Array::from(a.1.clone())), + Arc::new(Date32Array::from(b.1.clone())), + Arc::new(Date32Array::from(c.1.clone())), + ], + ) + .unwrap(); + + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +fn build_date64_table( + a: (&str, &Vec<i64>), + b: (&str, &Vec<i64>), + c: (&str, &Vec<i64>), +) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Date64, false), + Field::new(b.0, DataType::Date64, false), + Field::new(c.0, DataType::Date64, false), + ]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Date64Array::from(a.1.clone())), + Arc::new(Date64Array::from(b.1.clone())), + Arc::new(Date64Array::from(c.1.clone())), + ], + ) + .unwrap(); + + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +fn build_binary_table( + a: (&str, &Vec<&[u8]>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), +) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Binary, false), + Field::new(b.0, DataType::Int32, false), + Field::new(c.0, DataType::Int32, false), + ]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(BinaryArray::from(a.1.clone())), + Arc::new(Int32Array::from(b.1.clone())), + Arc::new(Int32Array::from(c.1.clone())), + ], + ) + .unwrap(); + + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +fn build_fixed_size_binary_table( + a: (&str, &Vec<&[u8]>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), +) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::FixedSizeBinary(3), false), + Field::new(b.0, DataType::Int32, false), + Field::new(c.0, DataType::Int32, false), + ]); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(FixedSizeBinaryArray::try_from_iter(a.1.iter().copied()).unwrap()), + Arc::new(Int32Array::from(b.1.clone())), + Arc::new(Int32Array::from(c.1.clone())), + ], + ) + .unwrap(); + + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +/// returns a table with 3 columns of i32 in memory +pub fn build_table_i32_nullable( + a: (&str, &Vec<Option<i32>>), + b: (&str, &Vec<Option<i32>>), + c: (&str, &Vec<Option<i32>>), +) -> Arc<dyn ExecutionPlan> { + let schema = Arc::new(Schema::new(vec![ + Field::new(a.0, DataType::Int32, true), + Field::new(b.0, DataType::Int32, true), + Field::new(c.0, DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(a.1.clone())), + Arc::new(Int32Array::from(b.1.clone())), + Arc::new(Int32Array::from(c.1.clone())), + ], + ) + .unwrap(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +pub fn build_table_two_cols( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), +) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32_two_cols(a, b); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +fn join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, +) -> Result<SortMergeJoinExec> { + let sort_options = vec![SortOptions::default(); on.len()]; + SortMergeJoinExec::try_new( + left, + right, + on, + None, + join_type, + sort_options, + NullEquality::NullEqualsNothing, + ) +} + +fn join_with_options( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, +) -> Result<SortMergeJoinExec> { + SortMergeJoinExec::try_new( + left, + right, + on, + None, + join_type, + sort_options, + null_equality, + ) +} + +fn join_with_filter( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: JoinFilter, + join_type: JoinType, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, +) -> Result<SortMergeJoinExec> { + SortMergeJoinExec::try_new( + left, + right, + on, + Some(filter), + join_type, + sort_options, + null_equality, + ) +} + +async fn join_collect( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, +) -> Result<(Vec<String>, Vec<RecordBatch>)> { + let sort_options = vec![SortOptions::default(); on.len()]; + join_collect_with_options( + left, + right, + on, + join_type, + sort_options, + NullEquality::NullEqualsNothing, + ) + .await +} + +async fn join_collect_with_filter( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: JoinFilter, + join_type: JoinType, +) -> Result<(Vec<String>, Vec<RecordBatch>)> { + let sort_options = vec![SortOptions::default(); on.len()]; + + let task_ctx = Arc::new(TaskContext::default()); + let join = join_with_filter( + left, + right, + on, + filter, + join_type, + sort_options, + NullEquality::NullEqualsNothing, + )?; + let columns = columns(&join.schema()); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + Ok((columns, batches)) +} + +async fn join_collect_with_options( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, + sort_options: Vec<SortOptions>, + null_equality: NullEquality, +) -> Result<(Vec<String>, Vec<RecordBatch>)> { + let task_ctx = Arc::new(TaskContext::default()); + let join = + join_with_options(left, right, on, join_type, sort_options, null_equality)?; + let columns = columns(&join.schema()); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + Ok((columns, batches)) +} + +async fn join_collect_batch_size_equals_two( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, +) -> Result<(Vec<String>, Vec<RecordBatch>)> { + let task_ctx = TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(2)); + let task_ctx = Arc::new(task_ctx); + let join = join(left, right, on, join_type)?; + let columns = columns(&join.schema()); + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + Ok((columns, batches)) +} + +#[tokio::test] +async fn join_inner_one() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 5]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 5 | 9 | 20 | 5 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_inner_two() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b2", &vec![1, 2, 2]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a1", &vec![1, 2, 3]), + ("b2", &vec![1, 2, 2]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let (_columns, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 1 | 7 | 1 | 1 | 70 | + | 2 | 2 | 8 | 2 | 2 | 80 | + | 2 | 2 | 9 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_inner_two_two() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 1, 2]), + ("b2", &vec![1, 1, 2]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a1", &vec![1, 1, 3]), + ("b2", &vec![1, 1, 2]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let (_columns, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 1 | 7 | 1 | 1 | 70 | + | 1 | 1 | 7 | 1 | 1 | 80 | + | 1 | 1 | 8 | 1 | 1 | 70 | + | 1 | 1 | 8 | 1 | 1 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_inner_with_nulls() -> Result<()> { + let left = build_table_i32_nullable( + ("a1", &vec![Some(1), Some(1), Some(2), Some(2)]), + ("b2", &vec![None, Some(1), Some(2), Some(2)]), // null in key field + ("c1", &vec![Some(1), None, Some(8), Some(9)]), // null in non-key field + ); + let right = build_table_i32_nullable( + ("a1", &vec![Some(1), Some(1), Some(2), Some(3)]), + ("b2", &vec![None, Some(1), Some(2), Some(2)]), + ("c2", &vec![Some(10), Some(70), Some(80), Some(90)]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 1 | | 1 | 1 | 70 | + | 2 | 2 | 8 | 2 | 2 | 80 | + | 2 | 2 | 9 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_inner_with_nulls_with_options() -> Result<()> { + let left = build_table_i32_nullable( + ("a1", &vec![Some(2), Some(2), Some(1), Some(1)]), + ("b2", &vec![Some(2), Some(2), Some(1), None]), // null in key field + ("c1", &vec![Some(9), Some(8), None, Some(1)]), // null in non-key field + ); + let right = build_table_i32_nullable( + ("a1", &vec![Some(3), Some(2), Some(1), Some(1)]), + ("b2", &vec![Some(2), Some(2), Some(1), None]), + ("c2", &vec![Some(90), Some(80), Some(70), Some(10)]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + let (_, batches) = join_collect_with_options( + left, + right, + on, + Inner, + vec![ + SortOptions { + descending: true, + nulls_first: false, + }; + 2 + ], + NullEquality::NullEqualsNull, + ) + .await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 2 | 2 | 9 | 2 | 2 | 80 | + | 2 | 2 | 8 | 2 | 2 | 80 | + | 1 | 1 | | 1 | 1 | 70 | + | 1 | | 1 | 1 | | 10 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_inner_output_two_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b2", &vec![1, 2, 2]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a1", &vec![1, 2, 3]), + ("b2", &vec![1, 2, 2]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b2", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect_batch_size_equals_two(left, right, on, Inner).await?; + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].num_rows(), 2); + assert_eq!(batches[1].num_rows(), 1); + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b2 | c1 | a1 | b2 | c2 | + +----+----+----+----+----+----+ + | 1 | 1 | 7 | 1 | 1 | 70 | + | 2 | 2 | 8 | 2 | 2 | 80 | + | 2 | 2 | 9 | 2 | 2 | 80 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_one() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Left).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_one() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Right).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +----+----+----+----+----+----+ + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | | | | 30 | 6 | 90 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_different_columns_count_with_filter() -> Result<()> { + // select * + // from t1 + // right join t2 on t1.b1 = t2.b1 and t1.a1 > t2.a2 + + let left = build_table( + ("a1", &vec![1, 21, 3]), // 21(t1.a1) > 20(t2.a2) + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + + let right = build_table_two_cols( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a1", 0)), + Operator::Gt, + Arc::new(Column::new("a2", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("a2", DataType::Int32, true), + ])), + ); + + let (_, batches) = join_collect_with_filter(left, right, on, filter, Right).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b1 | + +----+----+----+----+----+ + | | | | 10 | 4 | + | 21 | 5 | 8 | 20 | 5 | + | | | | 30 | 6 | + +----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_different_columns_count_with_filter() -> Result<()> { + // select * + // from t2 + // left join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 + + let left = build_table_two_cols( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the right + ); + + let right = build_table( + ("a1", &vec![1, 21, 3]), // 20(t2.a2) > 1(t1.a1) + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a2", 0)), + Operator::Gt, + Arc::new(Column::new("a1", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, true), + Field::new("a1", DataType::Int32, true), + ])), + ); + + let (_, batches) = join_collect_with_filter(left, right, on, filter, Left).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+ + | a2 | b1 | a1 | b1 | c1 | + +----+----+----+----+----+ + | 10 | 4 | 1 | 4 | 7 | + | 20 | 5 | | | | + | 30 | 6 | | | | + +----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_mark_different_columns_count_with_filter() -> Result<()> { + // select * + // from t2 + // left mark join t1 on t2.b1 = t1.b1 and t2.a2 > t1.a1 + + let left = build_table_two_cols( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the right + ); + + let right = build_table( + ("a1", &vec![1, 21, 3]), // 20(t2.a2) > 1(t1.a1) + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a2", 0)), + Operator::Gt, + Arc::new(Column::new("a1", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, true), + Field::new("a1", DataType::Int32, true), + ])), + ); + + let (_, batches) = + join_collect_with_filter(left, right, on, filter, LeftMark).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-------+ + | a2 | b1 | mark | + +----+----+-------+ + | 10 | 4 | true | + | 20 | 5 | false | + | 30 | 6 | false | + +----+----+-------+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_mark_different_columns_count_with_filter() -> Result<()> { + // select * + // from t1 + // right mark join t2 on t1.b1 = t2.b1 and t1.a1 > t2.a2 + + let left = build_table( + ("a1", &vec![1, 21, 3]), // 21(t1.a1) > 20(t2.a2) + ("b1", &vec![4, 5, 7]), + ("c1", &vec![7, 8, 9]), + ); + + let right = build_table_two_cols( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 6 does not exist on the left + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a1", 0)), + Operator::Gt, + Arc::new(Column::new("a2", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("a2", DataType::Int32, true), + ])), + ); + + let (_, batches) = + join_collect_with_filter(left, right, on, filter, RightMark).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-------+ + | a2 | b1 | mark | + +----+----+-------+ + | 10 | 4 | false | + | 20 | 5 | true | + | 30 | 6 | false | + +----+----+-------+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_full_one() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b2", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema()).unwrap()) as _, + Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Full).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 30 | 6 | 90 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | 3 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_anti() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2, 3, 5]), + ("b1", &vec![4, 5, 5, 7, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 8, 9, 11]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, LeftAnti).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 3 | 7 | 9 | + | 5 | 7 | 11 | + +----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_one_one() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b1", &vec![4, 5, 5]), + ("c1", &vec![7, 8, 8]), + ); + let right = build_table_two_cols(("a2", &vec![10, 20, 30]), ("b1", &vec![4, 5, 6])); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, RightAnti).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+ + | a2 | b1 | + +----+----+ + | 30 | 6 | + +----+----+ + "); + + let left2 = build_table( + ("a1", &vec![1, 2, 2]), + ("b1", &vec![4, 5, 5]), + ("c1", &vec![7, 8, 8]), + ); + let right2 = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left2.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right2.schema())?) as _, + )]; + + let (_, batches2) = join_collect(left2, right2, on, RightAnti).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches2), @r" + +----+----+----+ + | a2 | b1 | c2 | + +----+----+----+ + | 30 | 6 | 90 | + +----+----+----+ + "); + + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_two_two() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b1", &vec![4, 5, 5]), + ("c1", &vec![7, 8, 8]), + ); + let right = build_table_two_cols(("a2", &vec![10, 20, 30]), ("b1", &vec![4, 5, 6])); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect(left, right, on, RightAnti).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+ + | a2 | b1 | + +----+----+ + | 10 | 4 | + | 20 | 5 | + | 30 | 6 | + +----+----+ + "); + + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b1", &vec![4, 5, 5]), + ("c1", &vec![7, 8, 8]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect(left, right, on, RightAnti).await?; + let expected = [ + "+----+----+----+", + "| a2 | b1 | c2 |", + "+----+----+----+", + "| 10 | 4 | 70 |", + "| 20 | 5 | 80 |", + "| 30 | 6 | 90 |", + "+----+----+----+", + ]; + // The output order is important as SMJ preserves sortedness + assert_batches_eq!(expected, &batches); + + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_two_with_filter() -> Result<()> { + let left = build_table(("a1", &vec![1]), ("b1", &vec![10]), ("c1", &vec![30])); + let right = build_table(("a1", &vec![1]), ("b1", &vec![10]), ("c2", &vec![20])); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c2", 1)), + Operator::Gt, + Arc::new(Column::new("c1", 0)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, true), + Field::new("c2", DataType::Int32, true), + ])), + ); + let (_, batches) = + join_collect_with_filter(left, right, on, filter, RightAnti).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c2 | + +----+----+----+ + | 1 | 10 | 20 | + +----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_filtered_with_mismatched_columns() -> Result<()> { + let left = build_table_two_cols(("a1", &vec![31, 31]), ("b1", &vec![32, 33])); + let right = build_table( + ("a2", &vec![31, 31]), + ("b2", &vec![32, 35]), + ("c2", &vec![108, 109]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + ), + ]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b1", 0)), + Operator::LtEq, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let (_, batches) = + join_collect_with_filter(left, right, on, filter, RightAnti).await?; + + let expected = [ + "+----+----+-----+", + "| a2 | b2 | c2 |", + "+----+----+-----+", + "| 31 | 35 | 109 |", + "+----+----+-----+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_with_nulls() -> Result<()> { + let left = build_table_i32_nullable( + ("a1", &vec![Some(0), Some(1), Some(2), Some(2), Some(3)]), + ("b1", &vec![Some(3), Some(4), Some(5), None, Some(6)]), + ("c2", &vec![Some(60), None, Some(80), Some(85), Some(90)]), + ); + let right = build_table_i32_nullable( + ("a1", &vec![Some(1), Some(2), Some(2), Some(3)]), + ("b1", &vec![Some(4), Some(5), None, Some(6)]), // null in key field + ("c2", &vec![Some(7), Some(8), Some(8), None]), // null in non-key field + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect(left, right, on, RightAnti).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c2 | + +----+----+----+ + | 2 | | 8 | + +----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_with_nulls_with_options() -> Result<()> { + let left = build_table_i32_nullable( + ("a1", &vec![Some(1), Some(2), Some(1), Some(0), Some(2)]), + ("b1", &vec![Some(4), Some(5), Some(5), None, Some(5)]), + ("c1", &vec![Some(7), Some(8), Some(8), Some(60), None]), + ); + let right = build_table_i32_nullable( + ("a1", &vec![Some(3), Some(2), Some(2), Some(1)]), + ("b1", &vec![None, Some(5), Some(5), Some(4)]), // null in key field + ("c2", &vec![Some(9), None, Some(8), Some(7)]), // null in non-key field + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect_with_options( + left, + right, + on, + RightAnti, + vec![ + SortOptions { + descending: true, + nulls_first: false, + }; + 2 + ], + NullEquality::NullEqualsNull, + ) + .await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c2 | + +----+----+----+ + | 3 | | 9 | + | 2 | 5 | | + | 2 | 5 | 8 | + +----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_anti_output_two_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2]), + ("b1", &vec![4, 5, 5]), + ("c1", &vec![7, 8, 8]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = + join_collect_batch_size_equals_two(left, right, on, LeftAnti).await?; + // BitwiseSortMergeJoinStream uses a coalescer, so batch boundaries differ + // from the old stream. Only assert data correctness. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 3); + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 4 | 7 | + | 2 | 5 | 8 | + | 2 | 5 | 8 | + +----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_semi() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), // 5 is double on the right + ("c2", &vec![70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, LeftSemi).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+ + | a1 | b1 | c1 | + +----+----+----+ + | 1 | 4 | 7 | + | 2 | 5 | 8 | + | 2 | 5 | 8 | + +----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_semi_one() -> Result<()> { + let left = build_table( + ("a1", &vec![10, 20, 30, 40]), + ("b1", &vec![4, 5, 5, 6]), + ("c1", &vec![70, 80, 90, 100]), + ); + let right = build_table( + ("a2", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 7]), + ("c2", &vec![7, 8, 8, 9]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, RightSemi).await?; + let expected = [ + "+----+----+----+", + "| a2 | b1 | c2 |", + "+----+----+----+", + "| 1 | 4 | 7 |", + "| 2 | 5 | 8 |", + "| 2 | 5 | 8 |", + "+----+----+----+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_right_semi_two() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 6]), + ("c1", &vec![70, 80, 90, 100]), + ); + let right = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 7]), + ("c2", &vec![7, 8, 8, 9]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect(left, right, on, RightSemi).await?; + let expected = [ + "+----+----+----+", + "| a1 | b1 | c2 |", + "+----+----+----+", + "| 1 | 4 | 7 |", + "| 2 | 5 | 8 |", + "| 2 | 5 | 8 |", + "+----+----+----+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_right_semi_two_with_filter() -> Result<()> { + let left = build_table(("a1", &vec![1]), ("b1", &vec![10]), ("c1", &vec![30])); + let right = build_table(("a1", &vec![1]), ("b1", &vec![10]), ("c2", &vec![20])); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c2", 1)), + Operator::Lt, + Arc::new(Column::new("c1", 0)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, true), + Field::new("c2", DataType::Int32, true), + ])), + ); + let (_, batches) = + join_collect_with_filter(left, right, on, filter, RightSemi).await?; + let expected = [ + "+----+----+----+", + "| a1 | b1 | c2 |", + "+----+----+----+", + "| 1 | 10 | 20 |", + "+----+----+----+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_right_semi_with_nulls() -> Result<()> { + let left = build_table_i32_nullable( + ("a1", &vec![Some(0), Some(1), Some(2), Some(2), Some(3)]), + ("b1", &vec![Some(3), Some(4), Some(5), None, Some(6)]), + ("c2", &vec![Some(60), None, Some(80), Some(85), Some(90)]), + ); + let right = build_table_i32_nullable( + ("a1", &vec![Some(1), Some(2), Some(2), Some(3)]), + ("b1", &vec![Some(4), Some(5), None, Some(6)]), // null in key field + ("c2", &vec![Some(7), Some(8), Some(8), None]), // null in non-key field + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect(left, right, on, RightSemi).await?; + let expected = [ + "+----+----+----+", + "| a1 | b1 | c2 |", + "+----+----+----+", + "| 1 | 4 | 7 |", + "| 2 | 5 | 8 |", + "| 3 | 6 | |", + "+----+----+----+", + ]; + // The output order is important as SMJ preserves sortedness + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_right_semi_with_nulls_with_options() -> Result<()> { + let left = build_table_i32_nullable( + ("a1", &vec![Some(3), Some(2), Some(1), Some(0), Some(2)]), + ("b1", &vec![None, Some(5), Some(4), None, Some(5)]), + ("c2", &vec![Some(90), Some(80), Some(70), Some(60), None]), + ); + let right = build_table_i32_nullable( + ("a1", &vec![Some(3), Some(2), Some(2), Some(1)]), + ("b1", &vec![None, Some(5), Some(5), Some(4)]), // null in key field + ("c2", &vec![Some(9), None, Some(8), Some(7)]), // null in non-key field + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = join_collect_with_options( + left, + right, + on, + RightSemi, + vec![ + SortOptions { + descending: true, + nulls_first: false, + }; + 2 + ], + NullEquality::NullEqualsNull, + ) + .await?; + + let expected = [ + "+----+----+----+", + "| a1 | b1 | c2 |", + "+----+----+----+", + "| 3 | | 9 |", + "| 2 | 5 | |", + "| 2 | 5 | 8 |", + "| 1 | 4 | 7 |", + "+----+----+----+", + ]; + // The output order is important as SMJ preserves sortedness + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_right_semi_output_two_batches() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 6]), + ("c1", &vec![70, 80, 90, 100]), + ); + let right = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 7]), + ("c2", &vec![7, 8, 8, 9]), + ); + let on = vec![ + ( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + ), + ( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + ), + ]; + + let (_, batches) = + join_collect_batch_size_equals_two(left, right, on, RightSemi).await?; + let expected = [ + "+----+----+----+", + "| a1 | b1 | c2 |", + "+----+----+----+", + "| 1 | 4 | 7 |", + "| 2 | 5 | 8 |", + "| 2 | 5 | 8 |", + "+----+----+----+", + ]; + // BitwiseSortMergeJoinStream uses a coalescer, so batch boundaries differ + // from the old stream. Only assert data correctness. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 3); + assert_batches_eq!(expected, &batches); + Ok(()) +} + +#[tokio::test] +async fn join_left_mark() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40]), + ("b1", &vec![4, 4, 5, 6]), // 5 is double on the right + ("c2", &vec![60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, LeftMark).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+-------+ + | a1 | b1 | c1 | mark | + +----+----+----+-------+ + | 1 | 4 | 7 | true | + | 2 | 5 | 8 | true | + | 2 | 5 | 8 | true | + | 3 | 7 | 9 | false | + +----+----+----+-------+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_mark() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 2, 3]), + ("b1", &vec![4, 5, 5, 7]), // 7 does not exist on the right + ("c1", &vec![7, 8, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40]), + ("b1", &vec![4, 4, 5, 6]), // 5 is double on the left + ("c2", &vec![60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, RightMark).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+-------+ + | a2 | b1 | c2 | mark | + +----+----+----+-------+ + | 10 | 4 | 60 | true | + | 20 | 4 | 70 | true | + | 30 | 5 | 80 | true | + | 40 | 6 | 90 | false | + +----+----+----+-------+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_with_duplicated_column_names() -> Result<()> { + let left = build_table( + ("a", &vec![1, 2, 3]), + ("b", &vec![4, 5, 7]), + ("c", &vec![7, 8, 9]), + ); + let right = build_table( + ("a", &vec![10, 20, 30]), + ("b", &vec![1, 2, 7]), + ("c", &vec![70, 80, 90]), + ); + let on = vec![( + // join on a=b so there are duplicate column names on unjoined columns + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +---+---+---+----+---+----+ + | a | b | c | a | b | c | + +---+---+---+----+---+----+ + | 1 | 4 | 7 | 10 | 1 | 70 | + | 2 | 5 | 8 | 20 | 2 | 80 | + +---+---+---+----+---+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_date32() -> Result<()> { + let left = build_date_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![19107, 19108, 19108]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_date_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![19107, 19108, 19109]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +------------+------------+------------+------------+------------+------------+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +------------+------------+------------+------------+------------+------------+ + | 1970-01-02 | 2022-04-25 | 1970-01-08 | 1970-01-11 | 2022-04-25 | 1970-03-12 | + | 1970-01-03 | 2022-04-26 | 1970-01-09 | 1970-01-21 | 2022-04-26 | 1970-03-22 | + | 1970-01-04 | 2022-04-26 | 1970-01-10 | 1970-01-21 | 2022-04-26 | 1970-03-22 | + +------------+------------+------------+------------+------------+------------+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_date64() -> Result<()> { + let left = build_date64_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![1650703441000, 1650903441000, 1650903441000]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_date64_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![1650703441000, 1650503441000, 1650903441000]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + | a1 | b1 | c1 | a2 | b1 | c2 | + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + | 1970-01-01T00:00:00.001 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.007 | 1970-01-01T00:00:00.010 | 2022-04-23T08:44:01 | 1970-01-01T00:00:00.070 | + | 1970-01-01T00:00:00.002 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.008 | 1970-01-01T00:00:00.030 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.090 | + | 1970-01-01T00:00:00.003 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.009 | 1970-01-01T00:00:00.030 | 2022-04-25T16:17:21 | 1970-01-01T00:00:00.090 | + +-------------------------+---------------------+-------------------------+-------------------------+---------------------+-------------------------+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_binary() -> Result<()> { + let left = build_binary_table( + ( + "a1", + &vec![ + &[0xc0, 0xff, 0xee], + &[0xde, 0xca, 0xde], + &[0xfa, 0xca, 0xde], + ], + ), + ("b1", &vec![5, 10, 15]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_binary_table( + ( + "a1", + &vec![ + &[0xc0, 0xff, 0xee], + &[0xde, 0xca, 0xde], + &[0xfa, 0xca, 0xde], + ], + ), + ("b2", &vec![105, 110, 115]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +--------+----+----+--------+-----+----+ + | a1 | b1 | c1 | a1 | b2 | c2 | + +--------+----+----+--------+-----+----+ + | c0ffee | 5 | 7 | c0ffee | 105 | 70 | + | decade | 10 | 8 | decade | 110 | 80 | + | facade | 15 | 9 | facade | 115 | 90 | + +--------+----+----+--------+-----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_fixed_size_binary() -> Result<()> { + let left = build_fixed_size_binary_table( + ( + "a1", + &vec![ + &[0xc0, 0xff, 0xee], + &[0xde, 0xca, 0xde], + &[0xfa, 0xca, 0xde], + ], + ), + ("b1", &vec![5, 10, 15]), // this has a repetition + ("c1", &vec![7, 8, 9]), + ); + let right = build_fixed_size_binary_table( + ( + "a1", + &vec![ + &[0xc0, 0xff, 0xee], + &[0xde, 0xca, 0xde], + &[0xfa, 0xca, 0xde], + ], + ), + ("b2", &vec![105, 110, 115]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Inner).await?; + + // The output order is important as SMJ preserves sortedness + assert_snapshot!(batches_to_string(&batches), @r" + +--------+----+----+--------+-----+----+ + | a1 | b1 | c1 | a1 | b2 | c2 | + +--------+----+----+--------+-----+----+ + | c0ffee | 5 | 7 | c0ffee | 105 | 70 | + | decade | 10 | 8 | decade | 110 | 80 | + | facade | 15 | 9 | facade | 115 | 90 | + +--------+----+----+--------+-----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_sort_order() -> Result<()> { + let left = build_table( + ("a1", &vec![0, 1, 2, 3, 4, 5]), + ("b1", &vec![3, 4, 5, 6, 6, 7]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![0, 10, 20, 30, 40]), + ("b2", &vec![2, 4, 6, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Left).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 0 | 3 | 4 | | | | + | 1 | 4 | 5 | 10 | 4 | 60 | + | 2 | 5 | 6 | | | | + | 3 | 6 | 7 | 20 | 6 | 70 | + | 3 | 6 | 7 | 30 | 6 | 80 | + | 4 | 6 | 8 | 20 | 6 | 70 | + | 4 | 6 | 8 | 30 | 6 | 80 | + | 5 | 7 | 9 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_sort_order() -> Result<()> { + let left = build_table( + ("a1", &vec![0, 1, 2, 3]), + ("b1", &vec![3, 4, 5, 7]), + ("c1", &vec![6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![0, 10, 20, 30]), + ("b2", &vec![2, 4, 5, 6]), + ("c2", &vec![60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Right).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 0 | 2 | 60 | + | 1 | 4 | 7 | 10 | 4 | 70 | + | 2 | 5 | 8 | 20 | 5 | 80 | + | | | | 30 | 6 | 90 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_left_multiple_batches() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1, 2]), + ("b1", &vec![3, 4, 5]), + ("c1", &vec![4, 5, 6]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![3, 4, 5, 6]), + ("b1", &vec![6, 6, 7, 9]), + ("c1", &vec![7, 8, 9, 9]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10, 20]), + ("b2", &vec![2, 4, 6]), + ("c2", &vec![50, 60, 70]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![30, 40]), + ("b2", &vec![6, 8]), + ("c2", &vec![80, 90]), + ); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2]); + let right = build_table_from_batches(vec![right_batch_1, right_batch_2]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Left).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | 0 | 3 | 4 | | | | + | 1 | 4 | 5 | 10 | 4 | 60 | + | 2 | 5 | 6 | | | | + | 3 | 6 | 7 | 20 | 6 | 70 | + | 3 | 6 | 7 | 30 | 6 | 80 | + | 4 | 6 | 8 | 20 | 6 | 70 | + | 4 | 6 | 8 | 30 | 6 | 80 | + | 5 | 7 | 9 | | | | + | 6 | 9 | 9 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_right_multiple_batches() -> Result<()> { + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 1, 2]), + ("b2", &vec![3, 4, 5]), + ("c2", &vec![4, 5, 6]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![3, 4, 5, 6]), + ("b2", &vec![6, 6, 7, 9]), + ("c2", &vec![7, 8, 9, 9]), + ); + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 10, 20]), + ("b1", &vec![2, 4, 6]), + ("c1", &vec![50, 60, 70]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![30, 40]), + ("b1", &vec![6, 8]), + ("c1", &vec![80, 90]), + ); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2]); + let right = build_table_from_batches(vec![right_batch_1, right_batch_2]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Right).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 0 | 3 | 4 | + | 10 | 4 | 60 | 1 | 4 | 5 | + | | | | 2 | 5 | 6 | + | 20 | 6 | 70 | 3 | 6 | 7 | + | 30 | 6 | 80 | 3 | 6 | 7 | + | 20 | 6 | 70 | 4 | 6 | 8 | + | 30 | 6 | 80 | 4 | 6 | 8 | + | | | | 5 | 7 | 9 | + | | | | 6 | 9 | 9 | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn join_full_multiple_batches() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1, 2]), + ("b1", &vec![3, 4, 5]), + ("c1", &vec![4, 5, 6]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![3, 4, 5, 6]), + ("b1", &vec![6, 6, 7, 9]), + ("c1", &vec![7, 8, 9, 9]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10, 20]), + ("b2", &vec![2, 4, 6]), + ("c2", &vec![50, 60, 70]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![30, 40]), + ("b2", &vec![6, 8]), + ("c2", &vec![80, 90]), + ); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2]); + let right = build_table_from_batches(vec![right_batch_1, right_batch_2]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + + let (_, batches) = join_collect(left, right, on, Full).await?; + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+----+----+ + | a1 | b1 | c1 | a2 | b2 | c2 | + +----+----+----+----+----+----+ + | | | | 0 | 2 | 50 | + | | | | 40 | 8 | 90 | + | 0 | 3 | 4 | | | | + | 1 | 4 | 5 | 10 | 4 | 60 | + | 2 | 5 | 6 | | | | + | 3 | 6 | 7 | 20 | 6 | 70 | + | 3 | 6 | 7 | 30 | 6 | 80 | + | 4 | 6 | 8 | 20 | 6 | 70 | + | 4 | 6 | 8 | 30 | 6 | 80 | + | 5 | 7 | 9 | | | | + | 6 | 9 | 9 | | | | + +----+----+----+----+----+----+ + "); + Ok(()) +} + +/// Full outer join where the filter evaluates to NULL due to a nullable column. +/// NULL filter results must be treated as unmatched, not matched. +/// Reproducer for SPARK-43113. +#[tokio::test] +async fn join_full_null_filter_result() -> Result<()> { + // Left: (a, b) all non-null, sorted on a + let left = build_table_two_cols( + ("a1", &vec![1, 1, 2, 2, 3, 3]), + ("b1", &vec![1, 2, 1, 2, 1, 2]), + ); + + // Right: (a, b) with b nullable, sorted on a + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b2", DataType::Int32, true), + ])); + let right_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![None, Some(2)])), + ], + )?; + let right = + TestMemoryExec::try_new_exec(&[vec![right_batch]], right_schema, None).unwrap(); + + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + + // Filter: b1 < (b2 + 1) AND b1 < (a2 + 1) + // When b2 is NULL, (b2 + 1) is NULL, so b1 < NULL is NULL → unmatched. + let lit_1: PhysicalExprRef = Arc::new(Literal::new(ScalarValue::Int32(Some(1)))); + let b1_lt_b2_plus_1: PhysicalExprRef = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b1", 0)), + Operator::Lt, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("b2", 1)), + Operator::Plus, + Arc::clone(&lit_1), + )), + )); + let b1_lt_a2_plus_1: PhysicalExprRef = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b1", 0)), + Operator::Lt, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a2", 2)), + Operator::Plus, + Arc::clone(&lit_1), + )), + )); + let filter_expr: PhysicalExprRef = Arc::new(BinaryExpr::new( + b1_lt_b2_plus_1, + Operator::And, + b1_lt_a2_plus_1, + )); + + let filter = JoinFilter::new( + filter_expr, + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("b1", DataType::Int32, true), + Field::new("b2", DataType::Int32, true), + Field::new("a2", DataType::Int32, true), + ])), + ); + + let (_, batches) = join_collect_with_filter(left, right, on, filter, Full).await?; + + // r=(1,NULL): b2 is NULL → b1 < (NULL+1) is NULL → all a=1 rows unmatched + // r=(2,2): b1 < 3 AND b1 < 3 → both l=(2,1) and l=(2,2) match + // l=(3,*): no right row with a=3 → unmatched + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b2 | + +----+----+----+----+ + | | | 1 | | + | 1 | 1 | | | + | 1 | 2 | | | + | 2 | 1 | 2 | 2 | + | 2 | 2 | 2 | 2 | + | 3 | 1 | | | + | 3 | 2 | | | + +----+----+----+----+ + "); + Ok(()) +} + +#[tokio::test] +async fn overallocation_single_batch_no_spill() -> Result<()> { + let left = build_table( + ("a1", &vec![0, 1, 2, 3, 4, 5]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![0, 10, 20, 30, 40]), + ("b2", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let join_types = vec![ + // Semi/anti/mark joins use BitwiseSortMergeJoinStream which only tracks + // inner key buffer memory; tested in bitwise_sort_merge_join/tests.rs. + Inner, Left, Right, Full, + ]; + + // Disable DiskManager to prevent spilling + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let session_config = SessionConfig::default().with_batch_size(50); + + for join_type in join_types { + let task_ctx = TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)); + let task_ctx = Arc::new(task_ctx); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let err = common::collect(stream).await.unwrap_err(); + + assert_contains!(err.to_string(), "Failed to allocate additional"); + assert_contains!(err.to_string(), "SMJStream[0]"); + assert_contains!(err.to_string(), "Disk spilling disabled"); + assert!(join.metrics().is_some()); + assert_eq!(join.metrics().unwrap().spill_count(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_bytes(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_rows(), Some(0)); + } + + Ok(()) +} + +#[tokio::test] +async fn overallocation_multi_batch_no_spill() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![1, 1]), + ("c1", &vec![6, 7]), + ); + let left_batch_3 = build_table_i32( + ("a1", &vec![4, 5]), + ("b1", &vec![1, 1]), + ("c1", &vec![8, 9]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10]), + ("b2", &vec![1, 1]), + ("c2", &vec![50, 60]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![20, 30]), + ("b2", &vec![1, 1]), + ("c2", &vec![70, 80]), + ); + let right_batch_3 = + build_table_i32(("a2", &vec![40]), ("b2", &vec![1]), ("c2", &vec![90])); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2, left_batch_3]); + let right = + build_table_from_batches(vec![right_batch_1, right_batch_2, right_batch_3]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let join_types = vec![ + // Semi/anti/mark joins use BitwiseSortMergeJoinStream which only tracks + // inner key buffer memory; tested in bitwise_sort_merge_join/tests.rs. + Inner, Left, Right, Full, + ]; + + // Disable DiskManager to prevent spilling + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + let session_config = SessionConfig::default().with_batch_size(50); + + for join_type in join_types { + let task_ctx = TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)); + let task_ctx = Arc::new(task_ctx); + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let err = common::collect(stream).await.unwrap_err(); + + assert_contains!(err.to_string(), "Failed to allocate additional"); + assert_contains!(err.to_string(), "SMJStream[0]"); + assert_contains!(err.to_string(), "Disk spilling disabled"); + assert!(join.metrics().is_some()); + assert_eq!(join.metrics().unwrap().spill_count(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_bytes(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_rows(), Some(0)); + } + + Ok(()) +} + +#[tokio::test] +async fn overallocation_single_batch_spill() -> Result<()> { + let left = build_table( + ("a1", &vec![0, 1, 2, 3, 4, 5]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![0, 10, 20, 30, 40]), + ("b2", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let join_types = [ + // Semi/anti/mark joins use BitwiseSortMergeJoinStream which only tracks + // inner key buffer memory; tested in bitwise_sort_merge_join/tests.rs. + Inner, Left, Right, Full, + ]; + + // Enable DiskManager to allow spilling + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + for batch_size in [1, 50] { + let session_config = SessionConfig::default().with_batch_size(batch_size); + + for join_type in &join_types { + let task_ctx = TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)); + let task_ctx = Arc::new(task_ctx); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let spilled_join_result = common::collect(stream).await.unwrap(); + + assert!(join.metrics().is_some()); + assert!(join.metrics().unwrap().spill_count().unwrap() > 0); + assert!(join.metrics().unwrap().spilled_bytes().unwrap() > 0); + assert!(join.metrics().unwrap().spilled_rows().unwrap() > 0); + + // Run the test with no spill configuration as + let task_ctx_no_spill = + TaskContext::default().with_session_config(session_config.clone()); + let task_ctx_no_spill = Arc::new(task_ctx_no_spill); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx_no_spill)?; + let no_spilled_join_result = common::collect(stream).await.unwrap(); + + assert!(join.metrics().is_some()); + assert_eq!(join.metrics().unwrap().spill_count(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_bytes(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_rows(), Some(0)); + // Compare spilled and non spilled data to check spill logic doesn't corrupt the data + assert_eq!(spilled_join_result, no_spilled_join_result); + } + } + + Ok(()) +} + +#[tokio::test] +async fn overallocation_multi_batch_spill() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![1, 1]), + ("c1", &vec![6, 7]), + ); + let left_batch_3 = build_table_i32( + ("a1", &vec![4, 5]), + ("b1", &vec![1, 1]), + ("c1", &vec![8, 9]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10]), + ("b2", &vec![1, 1]), + ("c2", &vec![50, 60]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![20, 30]), + ("b2", &vec![1, 1]), + ("c2", &vec![70, 80]), + ); + let right_batch_3 = + build_table_i32(("a2", &vec![40]), ("b2", &vec![1]), ("c2", &vec![90])); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2, left_batch_3]); + let right = + build_table_from_batches(vec![right_batch_1, right_batch_2, right_batch_3]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let join_types = [ + // Semi/anti/mark joins use BitwiseSortMergeJoinStream which only tracks + // inner key buffer memory; tested in bitwise_sort_merge_join/tests.rs. + Inner, Left, Right, Full, + ]; + + // Enable DiskManager to allow spilling + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(500, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + for batch_size in [1, 50] { + let session_config = SessionConfig::default().with_batch_size(batch_size); + + for join_type in &join_types { + let task_ctx = TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)); + let task_ctx = Arc::new(task_ctx); + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let spilled_join_result = common::collect(stream).await.unwrap(); + assert!(join.metrics().is_some()); + assert!(join.metrics().unwrap().spill_count().unwrap() > 0); + assert!(join.metrics().unwrap().spilled_bytes().unwrap() > 0); + assert!(join.metrics().unwrap().spilled_rows().unwrap() > 0); + + // For Full joins, get_required_batch_indices extends 0..batches.len(), so + // poll_spilled_batches can restore all spilled batches at once via infallible + // grow(). Verify accounting tracked the transient spike and cleaned up. + let peak_mem = join + .metrics() + .and_then(|m| m.sum_by_name("peak_mem_used")) + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + peak_mem > 0, + "peak_mem_used should be > 0 for {join_type:?} batch_size={batch_size}" + ); + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "memory should be fully released after {join_type:?} completes + (batch_size={batch_size}): infallible grow during restore must be balanced" + ); + // Run the test with no spill configuration as + let task_ctx_no_spill = + TaskContext::default().with_session_config(session_config.clone()); + let task_ctx_no_spill = Arc::new(task_ctx_no_spill); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx_no_spill)?; + let no_spilled_join_result = common::collect(stream).await.unwrap(); + + assert!(join.metrics().is_some()); + assert_eq!(join.metrics().unwrap().spill_count(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_bytes(), Some(0)); + assert_eq!(join.metrics().unwrap().spilled_rows(), Some(0)); + // Compare spilled and non spilled data to check spill logic doesn't corrupt the data + assert_eq!(spilled_join_result, no_spilled_join_result); + } + } + + Ok(()) +} + +/// Verifies that `peak_mem_used` reflects join_arrays memory on the spill path. +/// +/// Uses a memory limit smaller than a single batch's `size_estimation` so that +/// every batch spills — the `Ok` arm of `allocate_reservation` is never hit. +/// Before the fix, `peak_mem_used` would stay 0 because `set_max` was only +/// called in the `Ok` arm. After the fix, the spill path calls +/// `grow(join_arrays_mem)` + `set_max`, so `peak_mem_used > 0`. +#[tokio::test] +async fn spill_join_arrays_memory_accounting() -> Result<()> { + use arrow::array::Array; + + let left_batch = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let size_estimation = left_batch.get_array_memory_size() + + Int32Array::from(vec![1, 1]).get_array_memory_size() + + 2usize.next_power_of_two() * size_of::<usize>() + + size_of::<std::ops::Range<usize>>() + + size_of::<usize>(); + let join_arrays_mem = Int32Array::from(vec![1, 1]).get_array_memory_size(); + + // Memory limit: too small for a full batch, large enough for join_arrays. + // Every batch hits the Err arm → spills → grow(join_arrays_mem). + let memory_limit = (size_estimation + join_arrays_mem) / 2; + assert!( + memory_limit < size_estimation && memory_limit > join_arrays_mem, + "limit {memory_limit} must be between join_arrays_mem {join_arrays_mem} \ + and size_estimation {size_estimation}" + ); + + let left_batches: Vec<RecordBatch> = (0..4) + .map(|i| { + build_table_i32( + ("a1", &vec![i * 2, i * 2 + 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![100 + i, 101 + i]), + ) + }) + .collect(); + let left = build_table_from_batches(left_batches); + + let right_batches: Vec<RecordBatch> = (0..2) + .map(|i| { + build_table_i32( + ("a2", &vec![i * 2, i * 2 + 1]), + ("b2", &vec![1, 1]), + ("c2", &vec![200 + i, 201 + i]), + ) + }) + .collect(); + let right = build_table_from_batches(right_batches); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let session_config = SessionConfig::default().with_batch_size(50); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime)), + ); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Inner, + sort_options, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let result = common::collect(stream).await.unwrap(); + + assert!(!result.is_empty(), "Expected non-empty join result"); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur" + ); + + // Before the fix, peak_mem_used was 0 here because set_max was only + // called in the Ok arm of allocate_reservation, which is never reached + // when every batch spills. After the fix, the spill path calls + // grow(join_arrays_mem) + set_max unconditionally. + let peak_mem = metrics + .sum_by_name("peak_mem_used") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + peak_mem >= join_arrays_mem, + "peak_mem_used ({peak_mem}) should be >= join_arrays_mem ({join_arrays_mem})" + ); + + // All memory must be released (grow/shrink balanced, no underflow) + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "All memory should be released after join completes" + ); + + Ok(()) +} + +/// Test the no-headroom scenario: pool is so tight that even +/// join_arrays_mem exceeds the pool limit. With force-grow, the +/// reservation still tracks the join_arrays unconditionally so the +/// pool reflects actual memory usage. +#[tokio::test] +async fn spill_join_arrays_no_headroom() -> Result<()> { + use arrow::array::Array; + + let join_arrays_mem = Int32Array::from(vec![1, 1]).get_array_memory_size(); + + // Pool smaller than join_arrays_mem: try_grow(size_estimation) fails → spill. + // Force-grow(join_arrays_mem) succeeds unconditionally → reserved_amount > 0. + let memory_limit = join_arrays_mem / 2; + assert!( + memory_limit < join_arrays_mem, + "limit {memory_limit} must be smaller than join_arrays_mem {join_arrays_mem}" + ); + + let left_batches: Vec<RecordBatch> = (0..4) + .map(|i| { + build_table_i32( + ("a1", &vec![i * 2, i * 2 + 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![100 + i, 101 + i]), + ) + }) + .collect(); + let left = build_table_from_batches(left_batches); + + let right_batches: Vec<RecordBatch> = (0..2) + .map(|i| { + build_table_i32( + ("a2", &vec![i * 2, i * 2 + 1]), + ("b2", &vec![1, 1]), + ("c2", &vec![200 + i, 201 + i]), + ) + }) + .collect(); + let right = build_table_from_batches(right_batches); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let session_config = SessionConfig::default().with_batch_size(50); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime)), + ); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Inner, + sort_options, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let result = common::collect(stream).await.unwrap(); + + assert!(!result.is_empty(), "Expected non-empty join result"); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur" + ); + + // Force-grow means peak_mem_used is always tracked, even when pool is tight. + let peak_mem = metrics + .sum_by_name("peak_mem_used") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + peak_mem >= join_arrays_mem, + "peak_mem_used ({peak_mem}) should be >= join_arrays_mem ({join_arrays_mem})" + ); + + // Pool should be fully released (grow/shrink balanced) + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "All memory should be released after join completes" + ); + + Ok(()) +} + +/// Build a c1 < c2 filter on the third column of each side. +fn build_c1_lt_c2_filter(left_schema: &Schema, right_schema: &Schema) -> JoinFilter { + JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Lt, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + left_schema + .field_with_name("c1") + .unwrap() + .clone() + .with_nullable(true), + right_schema + .field_with_name("c2") + .unwrap() + .clone() + .with_nullable(true), + ])), + ) +} + +#[tokio::test] +async fn spill_with_filter_deferred() -> Result<()> { + let left = build_table( + ("a1", &vec![0, 1, 2, 3, 4, 5]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![0, 10, 20, 30, 40]), + ("b2", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let filter = build_c1_lt_c2_filter(&left.schema(), &right.schema()); + + // Deferred filtering join types handled by the main MaterializingSortMergeJoinStream + let join_types = [Left, Right, Full]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + for batch_size in [1, 50] { + let session_config = SessionConfig::default().with_batch_size(batch_size); + + for join_type in &join_types { + // Run with spilling + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)), + ); + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + assert!(join.metrics().is_some()); + assert!( + join.metrics().unwrap().spill_count().unwrap() > 0, + "Expected spilling for {join_type:?} batch_size={batch_size}" + ); + + // Run without spilling + let task_ctx_no_spill = Arc::new( + TaskContext::default().with_session_config(session_config.clone()), + ); + let join_no_spill = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + let spilled_str = batches_to_sort_string(&spilled_result); + let no_spill_str = batches_to_sort_string(&no_spill_result); + assert_eq!( + spilled_str, no_spill_str, + "Spill vs no-spill mismatch for {join_type:?} batch_size={batch_size}" + ); + } + } + + Ok(()) +} + +#[tokio::test] +async fn spill_with_filter_multi_batch() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![1, 1]), + ("c1", &vec![6, 7]), + ); + let left_batch_3 = build_table_i32( + ("a1", &vec![4, 5]), + ("b1", &vec![1, 1]), + ("c1", &vec![8, 9]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10]), + ("b2", &vec![1, 1]), + ("c2", &vec![50, 60]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![20, 30]), + ("b2", &vec![1, 1]), + ("c2", &vec![70, 80]), + ); + let right_batch_3 = + build_table_i32(("a2", &vec![40]), ("b2", &vec![1]), ("c2", &vec![90])); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2, left_batch_3]); + let right = + build_table_from_batches(vec![right_batch_1, right_batch_2, right_batch_3]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let filter = build_c1_lt_c2_filter(&left.schema(), &right.schema()); + + let join_types = [Left, Right, Full]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(500, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + for batch_size in [1, 50] { + let session_config = SessionConfig::default().with_batch_size(batch_size); + + for join_type in &join_types { + // Run with spilling + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)), + ); + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + assert!(join.metrics().is_some()); + assert!( + join.metrics().unwrap().spill_count().unwrap() > 0, + "Expected spilling for {join_type:?} batch_size={batch_size}" + ); + + // Run without spilling + let task_ctx_no_spill = Arc::new( + TaskContext::default().with_session_config(session_config.clone()), + ); + let join_no_spill = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + *join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + let spilled_str = batches_to_sort_string(&spilled_result); + let no_spill_str = batches_to_sort_string(&no_spill_result); + assert_eq!( + spilled_str, no_spill_str, + "Spill vs no-spill mismatch for {join_type:?} batch_size={batch_size}" + ); + } + } + + Ok(()) +} + +/// FULL join where all buffered rows match on key but fail the filter. +/// Verifies produce_buffered_not_matched emits null-joined rows under spill. +#[tokio::test] +async fn spill_full_join_filter_not_matched() -> Result<()> { + // c1 values (100..105) are always > c2 values (1..5), so c1 < c2 always fails + let left = build_table( + ("a1", &vec![0, 1, 2, 3, 4]), + ("b1", &vec![1, 1, 1, 1, 1]), + ("c1", &vec![100, 101, 102, 103, 104]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40, 50]), + ("b2", &vec![1, 1, 1, 1, 1]), + ("c2", &vec![1, 2, 3, 4, 5]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let filter = build_c1_lt_c2_filter(&left.schema(), &right.schema()); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + for batch_size in [1, 50] { + let session_config = SessionConfig::default().with_batch_size(batch_size); + + // Run with spilling + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)), + ); + let join = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + Full, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + assert!( + join.metrics().unwrap().spill_count().unwrap() > 0, + "Expected spilling for FULL batch_size={batch_size}" + ); + + // Run without spilling + let task_ctx_no_spill = + Arc::new(TaskContext::default().with_session_config(session_config.clone())); + let join_no_spill = join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + Full, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + // All filter evaluations fail, so FULL join should produce: + // - 5 rows with left columns + null right columns (unmatched left) + // - 5 rows with null left columns + right columns (unmatched right) + let total_rows: usize = no_spill_result.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, 10, + "FULL join with all-failing filter should produce 10 rows, got {total_rows}" + ); + + let spilled_str = batches_to_sort_string(&spilled_result); + let no_spill_str = batches_to_sort_string(&no_spill_result); + assert_eq!( + spilled_str, no_spill_str, + "Spill vs no-spill mismatch for FULL join batch_size={batch_size}" + ); + } + + Ok(()) +} + +fn build_joined_record_batches() -> Result<JoinedRecordBatches> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ])); + + let mut batches = JoinedRecordBatches { + joined_batches: BatchCoalescer::new(Arc::clone(&schema), 8192), + filter_metadata: crate::joins::sort_merge_join::filter::FilterMetadata::new(), + }; + + // Insert already prejoined non-filtered rows + batches.joined_batches.push_batch(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![10, 10])), + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![11, 9])), + ], + )?)?; + + batches.joined_batches.push_batch(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![11])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![12])), + ], + )?)?; + + batches.joined_batches.push_batch(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![12, 12])), + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![11, 13])), + ], + )?)?; + + batches.joined_batches.push_batch(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![13])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![12])), + ], + )?)?; + + batches.joined_batches.push_batch(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![14, 14])), + Arc::new(Int32Array::from(vec![1, 1])), + Arc::new(Int32Array::from(vec![12, 11])), + ], + )?)?; + + let streamed_indices = vec![0, 0]; + batches + .filter_metadata + .batch_ids + .extend(vec![0; streamed_indices.len()]); + batches + .filter_metadata + .row_indices + .extend(&UInt64Array::from(streamed_indices)); + + let streamed_indices = vec![1]; + batches + .filter_metadata + .batch_ids + .extend(vec![0; streamed_indices.len()]); + batches + .filter_metadata + .row_indices + .extend(&UInt64Array::from(streamed_indices)); + + let streamed_indices = vec![0, 0]; + batches + .filter_metadata + .batch_ids + .extend(vec![1; streamed_indices.len()]); + batches + .filter_metadata + .row_indices + .extend(&UInt64Array::from(streamed_indices)); + + let streamed_indices = vec![0]; + batches + .filter_metadata + .batch_ids + .extend(vec![2; streamed_indices.len()]); + batches + .filter_metadata + .row_indices + .extend(&UInt64Array::from(streamed_indices)); + + let streamed_indices = vec![0, 0]; + batches + .filter_metadata + .batch_ids + .extend(vec![3; streamed_indices.len()]); + batches + .filter_metadata + .row_indices + .extend(&UInt64Array::from(streamed_indices)); + + batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![true, false])); + batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![true])); + batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![false, true])); + batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![false])); + batches + .filter_metadata + .filter_mask + .extend(&BooleanArray::from(vec![false, false])); + + Ok(batches) +} + +#[tokio::test] +async fn test_left_outer_join_filtered_mask() -> Result<()> { + let mut joined_batches = build_joined_record_batches()?; + let schema = joined_batches.joined_batches.schema(); + + let output = joined_batches.concat_batches(&schema)?; + let out_mask = joined_batches.filter_metadata.filter_mask.finish(); + let out_indices = joined_batches.filter_metadata.row_indices.finish(); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0]), + &[0usize], + &BooleanArray::from(vec![true]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![true, false, false, false, false, false, false, false]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0]), + &[0usize], + &BooleanArray::from(vec![false]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![false, false, false, false, false, false, false, false]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0, 0]), + &[0usize; 2], + &BooleanArray::from(vec![true, true]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![true, true, false, false, false, false, false, false]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0, 0, 0]), + &[0usize; 3], + &BooleanArray::from(vec![true, true, true]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![true, true, true, false, false, false, false, false]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0, 0, 0]), + &[0usize; 3], + &BooleanArray::from(vec![true, false, true]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![ + Some(true), + None, + Some(true), + Some(false), + Some(false), + Some(false), + Some(false), + Some(false) + ]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0, 0, 0]), + &[0usize; 3], + &BooleanArray::from(vec![false, false, true]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![ + None, + None, + Some(true), + Some(false), + Some(false), + Some(false), + Some(false), + Some(false) + ]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0, 0, 0]), + &[0usize; 3], + &BooleanArray::from(vec![false, true, true]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![ + None, + Some(true), + Some(true), + Some(false), + Some(false), + Some(false), + Some(false), + Some(false) + ]) + ); + + assert_eq!( + get_corrected_filter_mask( + Left, + &UInt64Array::from(vec![0, 0, 0]), + &[0usize; 3], + &BooleanArray::from(vec![false, false, false]), + output.num_rows() + ) + .unwrap(), + BooleanArray::from(vec![ + None, + None, + Some(false), + Some(false), + Some(false), + Some(false), + Some(false), + Some(false) + ]) + ); + + let corrected_mask = get_corrected_filter_mask( + Left, + &out_indices, + &joined_batches.filter_metadata.batch_ids, + &out_mask, + output.num_rows(), + ) + .unwrap(); + + assert_eq!( + corrected_mask, + BooleanArray::from(vec![ + Some(true), + None, + Some(true), + None, + Some(true), + Some(false), + None, + Some(false) + ]) + ); + + let filtered_rb = filter_record_batch(&output, &corrected_mask)?; + + assert_snapshot!(batches_to_string(&[filtered_rb]), @r" + +---+----+---+----+ + | a | b | x | y | + +---+----+---+----+ + | 1 | 10 | 1 | 11 | + | 1 | 11 | 1 | 12 | + | 1 | 12 | 1 | 13 | + +---+----+---+----+ + "); + + // output null rows + + let null_mask = arrow::compute::not(&corrected_mask)?; + assert_eq!( + null_mask, + BooleanArray::from(vec![ + Some(false), + None, + Some(false), + None, + Some(false), + Some(true), + None, + Some(true) + ]) + ); + + let null_joined_batch = filter_record_batch(&output, &null_mask)?; + + assert_snapshot!(batches_to_string(&[null_joined_batch]), @r" + +---+----+---+----+ + | a | b | x | y | + +---+----+---+----+ + | 1 | 13 | 1 | 12 | + | 1 | 14 | 1 | 11 | + +---+----+---+----+ + "); + Ok(()) +} + +#[test] +fn test_partition_statistics() -> Result<()> { + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use datafusion_common::stats::Precision; + + let left = build_table( + ("a1", &vec![1, 2, 3]), + ("b1", &vec![4, 5, 5]), + ("c1", &vec![7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30]), + ("b1", &vec![4, 5, 6]), + ("c2", &vec![70, 80, 90]), + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + // Test different join types to ensure partition_statistics works correctly for all + let join_types = vec![ + (Inner, 6), // left cols + right cols + (Left, 6), // left cols + right cols + (Right, 6), // left cols + right cols + (Full, 6), // left cols + right cols + (LeftSemi, 3), // only left cols + (LeftAnti, 3), // only left cols + (RightSemi, 3), // only right cols + (RightAnti, 3), // only right cols + ]; + + for (join_type, expected_cols) in join_types { + let join_exec = + join(Arc::clone(&left), Arc::clone(&right), on.clone(), join_type)?; + + // Test aggregate statistics (partition = None) + // Should return meaningful statistics computed from both inputs + let stats = + StatisticsContext::new().compute(&join_exec, &StatisticsArgs::new())?; + assert_eq!( + stats.column_statistics.len(), + expected_cols, + "Aggregate stats column count failed for {join_type:?}" + ); + // Verify that aggregate statistics have a meaningful num_rows (not Absent) + assert!( + stats.num_rows != Precision::Absent, + "Aggregate stats should have meaningful num_rows for {join_type:?}, got {:?}", + stats.num_rows + ); + + // Test partition-specific statistics (partition = Some(0)) + // The implementation correctly passes `partition` to children. + // Since the child TestMemoryExec returns unknown stats for specific partitions, + // the join output will also have Absent num_rows. This is expected behavior + // as the statistics depend on what the children can provide. + let partition_stats = StatisticsContext::new() + .compute(&join_exec, &StatisticsArgs::new().with_partition(Some(0)))?; + assert_eq!( + partition_stats.column_statistics.len(), + expected_cols, + "Partition stats column count failed for {join_type:?}" + ); + // When children return unknown stats, the join's partition stats will be Absent + assert!( + partition_stats.num_rows == Precision::Absent, + "Partition stats should have Absent num_rows when children return unknown for {join_type:?}, got {:?}", + partition_stats.num_rows + ); + } + + Ok(()) +} + +fn build_batches( + a: (&str, &[Vec<bool>]), + b: (&str, &[Vec<i32>]), + c: (&str, &[Vec<i32>]), +) -> (Vec<RecordBatch>, SchemaRef) { + assert_eq!(a.1.len(), b.1.len()); + let mut batches = vec![]; + + let schema = Arc::new(Schema::new(vec![ + Field::new(a.0, DataType::Boolean, false), + Field::new(b.0, DataType::Int32, false), + Field::new(c.0, DataType::Int32, false), + ])); + + for i in 0..a.1.len() { + batches.push( + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(BooleanArray::from(a.1[i].clone())), + Arc::new(Int32Array::from(b.1[i].clone())), + Arc::new(Int32Array::from(c.1[i].clone())), + ], + ) + .unwrap(), + ); + } + let schema = batches[0].schema(); + (batches, schema) +} + +fn build_batched_finish_barrier_table( + a: (&str, &[Vec<bool>]), + b: (&str, &[Vec<i32>]), + c: (&str, &[Vec<i32>]), +) -> (Arc<BarrierExec>, Arc<TestMemoryExec>) { + let (batches, schema) = build_batches(a, b, c); + + let memory_exec = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + ) + .unwrap(); + + let barrier_exec = Arc::new( + BarrierExec::new(vec![batches], schema) + .with_log(false) + .without_start_barrier() + .with_finish_barrier(), + ); + + (barrier_exec, memory_exec) +} + +/// Concat and sort batches by all the columns to make sure we can compare them with different join +fn prepare_record_batches_for_cmp(output: Vec<RecordBatch>) -> RecordBatch { + let output_batch = arrow::compute::concat_batches(output[0].schema_ref(), &output) + .expect("failed to concat batches"); + + // Sort on all columns to make sure we have a deterministic order for the assertion + let sort_columns = output_batch + .columns() + .iter() + .map(|c| SortColumn { + values: Arc::clone(c), + options: None, + }) + .collect::<Vec<_>>(); + + let sorted_columns = + arrow::compute::lexsort(&sort_columns, None).expect("failed to sort"); + + RecordBatch::try_new(output_batch.schema(), sorted_columns) + .expect("failed to create batch") +} + +#[expect(clippy::too_many_arguments)] +async fn join_get_stream_and_get_expected( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + oracle_left: Arc<dyn ExecutionPlan>, + oracle_right: Arc<dyn ExecutionPlan>, + on: JoinOn, + join_type: JoinType, + filter: Option<JoinFilter>, + batch_size: usize, +) -> Result<(SendableRecordBatchStream, RecordBatch)> { + let sort_options = vec![SortOptions::default(); on.len()]; + let null_equality = NullEquality::NullEqualsNothing; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::default().with_batch_size(batch_size)), + ); + + let expected_output = { + let oracle = HashJoinExec::try_new( + oracle_left, + oracle_right, + on.clone(), + filter.clone(), + &join_type, + None, + PartitionMode::Partitioned, + null_equality, + false, + )?; + + let stream = oracle.execute(0, Arc::clone(&task_ctx))?; + + let batches = common::collect(stream).await?; + + prepare_record_batches_for_cmp(batches) + }; + + let join = SortMergeJoinExec::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )?; + + let stream = join.execute(0, task_ctx)?; + + Ok((stream, expected_output)) +} + +fn generate_data_for_emit_early_test( + batch_size: usize, + number_of_batches: usize, + join_type: JoinType, +) -> ( + Arc<BarrierExec>, + Arc<BarrierExec>, + Arc<TestMemoryExec>, + Arc<TestMemoryExec>, +) { + let number_of_rows_per_batch = number_of_batches * batch_size; + // Prepare data + let left_a1 = (0..number_of_rows_per_batch as i32) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::<Vec<_>>()) + .collect::<Vec<_>>(); + let left_b1 = (0..1000000) + .filter(|item| { + match join_type { + LeftAnti | RightAnti => { + let remainder = item % (batch_size as i32); + + // Make sure to have one that match and one that don't + remainder == 0 || remainder == 1 + } + // Have at least 1 that is not matching + _ => item % batch_size as i32 != 0, + } + }) + .take(number_of_rows_per_batch) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::<Vec<_>>()) + .collect::<Vec<_>>(); + + let left_bool_col1 = left_a1 + .clone() + .into_iter() + .map(|b| { + b.into_iter() + // Mostly true but have some false that not overlap with the right column + .map(|a| a % (batch_size as i32) != (batch_size as i32) - 2) + .collect::<Vec<_>>() + }) + .collect::<Vec<_>>(); + + let (left, left_memory) = build_batched_finish_barrier_table( + ("bool_col1", left_bool_col1.as_slice()), + ("b1", left_b1.as_slice()), + ("a1", left_a1.as_slice()), + ); + + let right_a2 = (0..number_of_rows_per_batch as i32) + .map(|item| item * 11) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::<Vec<_>>()) + .collect::<Vec<_>>(); + let right_b1 = (0..1000000) + .filter(|item| { + match join_type { + LeftAnti | RightAnti => { + let remainder = item % (batch_size as i32); + + // Make sure to have one that match and one that don't + remainder == 1 || remainder == 2 + } + // Have at least 1 that is not matching + _ => item % batch_size as i32 != 1, + } + }) + .take(number_of_rows_per_batch) + .chunks(batch_size) + .into_iter() + .map(|chunk| chunk.collect::<Vec<_>>()) + .collect::<Vec<_>>(); + let right_bool_col2 = right_a2 + .clone() + .into_iter() + .map(|b| { + b.into_iter() + // Mostly true but have some false that not overlap with the left column + .map(|a| a % (batch_size as i32) != (batch_size as i32) - 1) + .collect::<Vec<_>>() + }) + .collect::<Vec<_>>(); + + let (right, right_memory) = build_batched_finish_barrier_table( + ("bool_col2", right_bool_col2.as_slice()), + ("b1", right_b1.as_slice()), + ("a2", right_a2.as_slice()), + ); + + (left, right, left_memory, right_memory) +} + +#[tokio::test] +async fn test_should_emit_early_when_have_enough_data_to_emit() -> Result<()> { + for with_filtering in [false, true] { + let join_types = vec![ + Inner, Left, Right, RightSemi, Full, LeftSemi, LeftAnti, LeftMark, RightMark, + ]; + const BATCH_SIZE: usize = 10; + for join_type in join_types { + for output_batch_size in [ + BATCH_SIZE / 3, + BATCH_SIZE / 2, + BATCH_SIZE, + BATCH_SIZE * 2, + BATCH_SIZE * 3, + ] { + // Make sure the number of batches is enough for all join type to emit some output + let number_of_batches = if output_batch_size <= BATCH_SIZE { + 100 + } else { + // Have enough batches + (output_batch_size * 100) / BATCH_SIZE + }; + + let (left, right, left_memory, right_memory) = + generate_data_for_emit_early_test( + BATCH_SIZE, + number_of_batches, + join_type, + ); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + let join_filter = if with_filtering { + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("bool_col1", 0)), + Operator::And, + Arc::new(Column::new("bool_col2", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("bool_col1", DataType::Boolean, true), + Field::new("bool_col2", DataType::Boolean, true), + ])), + ); + Some(filter) + } else { + None + }; + + // select * + // from t1 + // right join t2 on t1.b1 = t2.b1 and t1.bool_col1 AND t2.bool_col2 + let (mut output_stream, expected) = join_get_stream_and_get_expected( + Arc::clone(&left) as Arc<dyn ExecutionPlan>, + Arc::clone(&right) as Arc<dyn ExecutionPlan>, + left_memory as Arc<dyn ExecutionPlan>, + right_memory as Arc<dyn ExecutionPlan>, + on, + join_type, + join_filter, + output_batch_size, + ) + .await?; + + let (output_batched, output_batches_after_finish) = + consume_stream_until_finish_barrier_reached(left, right, &mut output_stream).await.unwrap_or_else(|e| panic!("Failed to consume stream for join type: '{join_type}' and with filtering '{with_filtering}': {e:?}")); + + // It should emit more than that, but we are being generous + // and to make sure the test pass for all + const MINIMUM_OUTPUT_BATCHES: usize = 5; + assert!( + MINIMUM_OUTPUT_BATCHES <= number_of_batches / 5, + "Make sure that the minimum output batches is realistic" + ); + // Test to make sure that we are not waiting for input to be fully consumed to emit some output + assert!( + output_batched.len() >= MINIMUM_OUTPUT_BATCHES, + "[Sort Merge Join {join_type}] Stream must have at least emit {} batches, but only got {} batches", + MINIMUM_OUTPUT_BATCHES, + output_batched.len() + ); + + // Just sanity test to make sure we are still producing valid output + { + let output = [output_batched, output_batches_after_finish].concat(); + let actual_prepared = prepare_record_batches_for_cmp(output); + + assert_eq!(actual_prepared.columns(), expected.columns()); + } + } + } + } + Ok(()) +} + +/// Polls the stream until both barriers are reached, +/// collecting the emitted batches along the way. +/// +/// If the stream is pending for too long (5s) without emitting any batches, +/// it panics to avoid hanging the test indefinitely. +/// +/// Note: The left and right BarrierExec might be the input of the output stream +async fn consume_stream_until_finish_barrier_reached( + left: Arc<BarrierExec>, + right: Arc<BarrierExec>, + output_stream: &mut SendableRecordBatchStream, +) -> Result<(Vec<RecordBatch>, Vec<RecordBatch>)> { + let mut switch_to_finish_barrier = false; + let mut output_batched = vec![]; + let mut after_finish_barrier_reached = vec![]; + let mut background_task = JoinSet::new(); + + let mut start_time_since_last_ready = Instant::now(); + loop { + let next_item = output_stream.next(); + + // Manual polling + let poll_output = futures::poll!(next_item); + + // Wake up the stream to make sure it makes progress + tokio::task::yield_now().await; + + match poll_output { + Poll::Ready(Some(Ok(batch))) => { + if batch.num_rows() == 0 { + return internal_err!("join stream should not emit empty batch"); + } + if switch_to_finish_barrier { + after_finish_barrier_reached.push(batch); + } else { + output_batched.push(batch); + } + start_time_since_last_ready = Instant::now(); + } + Poll::Ready(Some(Err(e))) => return Err(e), + Poll::Ready(None) if !switch_to_finish_barrier => { + unreachable!("Stream should not end before manually finishing it") + } + Poll::Ready(None) => { + break; + } + Poll::Pending => { + if right.is_finish_barrier_reached() + && left.is_finish_barrier_reached() + && !switch_to_finish_barrier + { + switch_to_finish_barrier = true; + + let right = Arc::clone(&right); + background_task.spawn(async move { + right.wait_finish().await; + }); + let left = Arc::clone(&left); + background_task.spawn(async move { + left.wait_finish().await; + }); + } + + // Make sure the test doesn't run forever + if start_time_since_last_ready.elapsed() > Duration::from_secs(5) { + return internal_err!( + "Stream should have emitted data by now, but it's still pending. Output batches so far: {}", + output_batched.len() + ); + } + } + } + } + + Ok((output_batched, after_finish_barrier_reached)) +} + +/// Exercises the multi-source interleave path in `materialize_right_columns`. +/// +/// When the right (buffered) side is split into many small batches with unique +/// keys, a single `freeze_streamed()` call references multiple `BufferedBatch`es. +/// This forces the `interleave` kernel instead of the single-source `take` path. +/// Without this test, the interleave path has zero coverage from unit tests +/// (fuzz tests use ~100 unique keys across 1000 rows, so all keys fit in one +/// buffered batch). +#[tokio::test] +async fn join_filtered_with_multiple_buffered_batches() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("val_l", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("val_r", DataType::Int32, false), + ])); + + // Left: single batch, keys 1..=6 + let left_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])), + Arc::new(Int32Array::from(vec![10, 20, 30, 40, 50, 60])), + ], + )?; + let left = build_table_from_batches(vec![left_batch]); + + // Right: one row per batch so each key lives in a separate BufferedBatch + let right_batches: Vec<RecordBatch> = (1..=6) + .map(|k| { + RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![k])), + Arc::new(Int32Array::from(vec![k * 100])), + ], + ) + .unwrap() + }) + .collect(); + let right = build_table_from_batches(right_batches); + + let on: JoinOn = vec![( + Arc::new(Column::new_with_schema("key", &left.schema())?) as _, + Arc::new(Column::new_with_schema("key", &right.schema())?) as _, + )]; + + // Filter: val_l + val_r < 350 — passes for keys 1-3, fails for 4-6 + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("val_l", 0)), + Operator::Plus, + Arc::new(Column::new("val_r", 1)), + )), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(350)))), + )), + vec![ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("val_l", DataType::Int32, true), + Field::new("val_r", DataType::Int32, true), + ])), + ); + + // Inner: only rows passing the filter + let (_, batches) = join_collect_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + Inner, + ) + .await?; + let result = batches_to_sort_string(&batches); + assert_snapshot!(result, @r" + +-----+-------+-----+-------+ + | key | val_l | key | val_r | + +-----+-------+-----+-------+ + | 1 | 10 | 1 | 100 | + | 2 | 20 | 2 | 200 | + | 3 | 30 | 3 | 300 | + +-----+-------+-----+-------+ + "); + + // Left: unmatched left rows get null right columns + let (_, batches) = join_collect_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + Left, + ) + .await?; + let result = batches_to_sort_string(&batches); + assert_snapshot!(result, @r" + +-----+-------+-----+-------+ + | key | val_l | key | val_r | + +-----+-------+-----+-------+ + | 1 | 10 | 1 | 100 | + | 2 | 20 | 2 | 200 | + | 3 | 30 | 3 | 300 | + | 4 | 40 | | | + | 5 | 50 | | | + | 6 | 60 | | | + +-----+-------+-----+-------+ + "); + + // Full: unmatched rows on both sides get null columns + let (_, batches) = join_collect_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + Full, + ) + .await?; + let result = batches_to_sort_string(&batches); + assert_snapshot!(result, @r" + +-----+-------+-----+-------+ + | key | val_l | key | val_r | + +-----+-------+-----+-------+ + | | | 4 | 400 | + | | | 5 | 500 | + | | | 6 | 600 | + | 1 | 10 | 1 | 100 | + | 2 | 20 | 2 | 200 | + | 3 | 30 | 3 | 300 | + | 4 | 40 | | | + | 5 | 50 | | | + | 6 | 60 | | | + +-----+-------+-----+-------+ + "); + + Ok(()) +} + +/// Returns the column names on the schema +fn columns(schema: &Schema) -> Vec<String> { + schema.fields().iter().map(|f| f.name().clone()).collect() +} + +// ==================== BitwiseSortMergeJoinStream direct tests ==================== +// +// These tests construct a BitwiseSortMergeJoinStream directly (bypassing exec) +// to exercise waiting on inputs and spill edge cases using PendingStream. + +/// Create test memory/spill resources for stream-level tests. +fn test_stream_resources( + inner_schema: SchemaRef, + metrics: &ExecutionPlanMetricsSet, +) -> ( + datafusion_execution::memory_pool::MemoryReservation, + SpillManager, + Arc<datafusion_execution::runtime_env::RuntimeEnv>, +) { + let ctx = TaskContext::default(); + let runtime_env = ctx.runtime_env(); + let reservation = MemoryConsumer::new("test").register(ctx.memory_pool()); + let spill_manager = SpillManager::new( + Arc::clone(&runtime_env), + SpillMetrics::new(metrics, 0), + inner_schema, + ); + (reservation, spill_manager, runtime_env) +} + +/// A RecordBatch stream that yields Poll::Pending once before delivering +/// each batch at a specified index. This simulates the behavior of +/// repartitioned tokio::sync::mpsc channels where data isn't immediately +/// available. +struct PendingStream { + batches: Vec<RecordBatch>, + index: usize, + /// If pending_before[i] is true, yield Pending once before delivering + /// the batch at index i. + pending_before: Vec<bool>, + /// True if we've already yielded Pending for the current index. + yielded_pending: bool, + schema: SchemaRef, +} + +impl PendingStream { + fn new(batches: Vec<RecordBatch>, pending_before: Vec<bool>) -> Self { + assert_eq!(batches.len(), pending_before.len()); + let schema = batches[0].schema(); + Self { + batches, + index: 0, + pending_before, + yielded_pending: false, + schema, + } + } +} + +impl Stream for PendingStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.index >= self.batches.len() { + return Poll::Ready(None); + } + if self.pending_before[self.index] && !self.yielded_pending { + self.yielded_pending = true; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + self.yielded_pending = false; + let batch = self.batches[self.index].clone(); + self.index += 1; + Poll::Ready(Some(Ok(batch))) + } +} + +impl RecordBatchStream for PendingStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Helper: collect all output from a BitwiseSortMergeJoinStream. +async fn collect_stream(stream: SendableRecordBatchStream) -> Result<Vec<RecordBatch>> { + common::collect(stream).await +} + +// ==================== join_time metric tests ==================== +// +// These verify that `join_time` measures only the join's own work: waiting +// for either child input or for the consumer to take an emitted batch must +// not be counted. + +/// Stream that sleeps `delay` before yielding each batch, to simulate a +/// slow input. +fn delayed_stream( + batches: Vec<RecordBatch>, + delay: Duration, +) -> SendableRecordBatchStream { + let schema = batches[0].schema(); + Box::pin(crate::stream::RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches.into_iter().map(Ok)).then(move |item| async move { + tokio::time::sleep(delay).await; + item + }), + )) +} + +/// Three 2-row batches with unique matching keys. +fn join_time_batches() -> Vec<RecordBatch> { + vec![ + build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 2]), + ("c1", &vec![7, 8]), + ), + build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![3, 4]), + ("c1", &vec![7, 8]), + ), + build_table_i32( + ("a1", &vec![4, 5]), + ("b1", &vec![5, 6]), + ("c1", &vec![7, 8]), + ), + ] +} + +/// Build a no-filter LeftSemi bitwise stream over the given input streams. +/// The small batch size makes each outer batch surface as its own output +/// batch, so a slow consumer test sees multiple emits. +fn join_time_test_join( + outer: SendableRecordBatchStream, + inner: SendableRecordBatchStream, +) -> (SendableRecordBatchStream, ExecutionPlanMetricsSet) { + let metrics = ExecutionPlanMetricsSet::new(); + let outer_schema = outer.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner.schema(), &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + outer_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + vec![Arc::new(Column::new("b1", 1)) as PhysicalExprRef], + vec![Arc::new(Column::new("b1", 1)) as PhysicalExprRef], + None, + LeftSemi, + 2, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + ) + .unwrap(); + (stream, metrics) +} + +fn join_time_of(metrics: &ExecutionPlanMetricsSet) -> Duration { + Duration::from_nanos( + metrics + .clone_inner() + .sum_by_name("join_time") + .map(|m| m.as_usize()) + .unwrap_or(0) as u64, + ) +} + +/// Run a join with the given injected `delay`, retrying with 4x the delay +/// (up to 3 attempts) when `join_time < delay` fails. +/// +/// This de-flakes the check without masking real bugs: a genuine exclusion +/// bug makes `join_time` absorb the injected waits, so it scales with the +/// delay and fails at every escalation level. Only a fixed-size disturbance +/// (e.g. the OS preempting the test thread while the join_time clock is +/// running) is filtered out, since it cannot grow 4x with the delay. +/// +/// `run` returns `(join_time, wall)` for one join execution. Deterministic +/// invariants (row counts, wall-time lower bounds) stay as asserts inside +/// `run` — deliberately: a panic there fails the test immediately without +/// retrying, since those cannot flake and escalation would only mask a real +/// bug. Likewise `Err` from `run` (join execution failure) propagates +/// immediately. Only the preemption-sensitive `join_time` check is retried. +async fn check_join_time_excluded<F, Fut>(mut run: F) -> Result<()> +where + F: FnMut(Duration) -> Fut, + Fut: Future<Output = Result<(Duration, Duration)>>, +{ + let mut delay = Duration::from_millis(50); + for attempt in 0..3 { + let (join_time, wall) = run(delay).await?; + if join_time < delay { + return Ok(()); + } + assert!( + attempt < 2, + "join_time ({join_time:?}) should be well below the injected \ + delay ({delay:?}) even after escalating retries; wall {wall:?}" + ); + delay *= 4; + } + unreachable!() +} + +/// join_time must not include time spent waiting for the outer input. +#[tokio::test] +async fn join_time_excludes_outer_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let outer = delayed_stream(join_time_batches(), delay); + let inner = delayed_stream(join_time_batches(), Duration::ZERO); + let (stream, metrics) = join_time_test_join(outer, inner); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all outer rows should match"); + assert!( + wall >= delay * 3, + "outer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time spent waiting for the inner input. +#[tokio::test] +async fn join_time_excludes_inner_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let outer = delayed_stream(join_time_batches(), Duration::ZERO); + let inner = delayed_stream(join_time_batches(), delay); + let (stream, metrics) = join_time_test_join(outer, inner); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all outer rows should match"); + assert!( + wall >= delay * 3, + "inner delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time the consumer spends holding an emitted +/// batch (the generator is suspended inside `emitter.emit` meanwhile). +#[tokio::test] +async fn join_time_excludes_consumer_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let outer = delayed_stream(join_time_batches(), Duration::ZERO); + let inner = delayed_stream(join_time_batches(), Duration::ZERO); + let (mut stream, metrics) = join_time_test_join(outer, inner); + + let start = Instant::now(); + let mut output_batches = 0u32; + while let Some(batch) = stream.next().await { + batch?; + output_batches += 1; + // Simulate a slow consumer between emitted batches. + tokio::time::sleep(delay).await; + } + let wall = start.elapsed(); + + assert!( + output_batches >= 3, + "expected multiple emitted batches, got {output_batches}" + ); + assert!( + wall >= delay * output_batches, + "consumer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// Three 2-row batches with unique matching keys, right-side column names. +fn join_time_batches_right() -> Vec<RecordBatch> { + vec![ + build_table_i32( + ("a2", &vec![0, 1]), + ("b2", &vec![1, 2]), + ("c2", &vec![7, 8]), + ), + build_table_i32( + ("a2", &vec![2, 3]), + ("b2", &vec![3, 4]), + ("c2", &vec![7, 8]), + ), + build_table_i32( + ("a2", &vec![4, 5]), + ("b2", &vec![5, 6]), + ("c2", &vec![7, 8]), + ), + ] +} + +/// Build a no-filter Inner materializing join over the given input streams. +/// The small batch size makes the output surface as multiple batches, so a +/// slow consumer test sees multiple emits. +fn materializing_join_time_test_join( + streamed: SendableRecordBatchStream, + buffered: SendableRecordBatchStream, +) -> (SendableRecordBatchStream, ExecutionPlanMetricsSet) { + use crate::joins::sort_merge_join::materializing_stream::MaterializingSortMergeJoinStream; + use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; + + let metrics = ExecutionPlanMetricsSet::new(); + let out_schema = Arc::new(Schema::new( + streamed + .schema() + .fields() + .iter() + .chain(buffered.schema().fields().iter()) + .map(|f| f.as_ref().clone()) + .collect::<Vec<_>>(), + )); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(buffered.schema(), &metrics); + let stream = MaterializingSortMergeJoinStream::try_new( + out_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + streamed, + buffered, + vec![Arc::new(Column::new("b1", 1)) as _], + vec![Arc::new(Column::new("b2", 1)) as _], + None, + Inner, + 2, + SortMergeJoinMetrics::new(0, &metrics), + reservation, + spill_manager, + runtime_env, + ) + .unwrap(); + (stream, metrics) +} + +/// join_time must not include time spent waiting for the streamed input. +#[tokio::test] +async fn materializing_join_time_excludes_streamed_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), delay); + let buffered = delayed_stream(join_time_batches_right(), Duration::ZERO); + let (stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all rows should match"); + assert!( + wall >= delay * 3, + "streamed delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time spent waiting for the buffered input. +#[tokio::test] +async fn materializing_join_time_excludes_buffered_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), Duration::ZERO); + let buffered = delayed_stream(join_time_batches_right(), delay); + let (stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all rows should match"); + assert!( + wall >= delay * 3, + "buffered delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time the consumer spends holding an emitted +/// batch (the generator is suspended inside `emitter.emit` meanwhile). +#[tokio::test] +async fn materializing_join_time_excludes_consumer_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), Duration::ZERO); + let buffered = delayed_stream(join_time_batches_right(), Duration::ZERO); + let (mut stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let mut output_batches = 0u32; + while let Some(batch) = stream.next().await { + batch?; + output_batches += 1; + // Simulate a slow consumer between emitted batches. + tokio::time::sleep(delay).await; + } + let wall = start.elapsed(); + + assert!( + output_batches >= 3, + "expected multiple emitted batches, got {output_batches}" + ); + assert!( + wall >= delay * output_batches, + "consumer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// An inner key group spanning multiple inner batches must survive the inner +/// input returning Pending mid-way: inner rows delivered before the Pending +/// still take part in the filter evaluation. +/// +/// Setup: +/// - Inner: 3 single-row batches, all with key=1, filter values c2=[10, 20, 30] +/// - Outer: 1 row, key=1, filter value c1=10 +/// - Filter: c1 == c2 (only first inner row c2=10 matches) +/// - Pending injected before 3rd inner batch +/// +/// Expected: outer row emitted (match via c2=10) +#[tokio::test] +async fn filter_buffer_pending_loses_inner_rows() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + + // Outer: 1 row, key=1, c1=10 + let outer_batch = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![1])), // join key + Arc::new(Int32Array::from(vec![10])), // filter value + ], + )?; + + // Inner: 3 single-row batches, key=1, c2=[10, 20, 30] + let inner_batch1 = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![100])), + Arc::new(Int32Array::from(vec![1])), // join key + Arc::new(Int32Array::from(vec![10])), // matches filter + ], + )?; + let inner_batch2 = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![200])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![20])), // doesn't match + ], + )?; + let inner_batch3 = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![300])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![30])), // doesn't match + ], + )?; + + let outer: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![outer_batch], + vec![false], // outer delivers immediately + )); + let inner: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![inner_batch1, inner_batch2, inner_batch3], + vec![false, false, true], // Pending before 3rd batch + )); + + // Filter: c1 == c2 + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Eq, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let on_outer: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + let on_inner: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + + let metrics = ExecutionPlanMetricsSet::new(); + let inner_schema = inner.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner_schema, &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + left_schema, // output schema = outer schema for semi + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + on_outer, + on_inner, + Some(filter), + LeftSemi, + 8192, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + )?; + + let batches = collect_stream(stream).await?; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 1, + "LeftSemi with filter: outer row should be emitted because \ + inner row c2=10 matches filter c1==c2. Got {total} rows." + ); + Ok(()) +} + +/// A matched outer key group spanning a batch boundary must survive the outer +/// input returning Pending at that boundary: the rows continuing the key group +/// still count as matched, even though the inner side has already advanced +/// past the key. +/// +/// Setup: +/// - Outer: 2 single-row batches, both with key=1 (key group spans boundary) +/// - Inner: 1 row with key=1 +/// - Pending injected on outer before 2nd batch +/// +/// Expected: both outer rows emitted +#[tokio::test] +async fn no_filter_boundary_pending_loses_outer_rows() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + + // Outer: 2 single-row batches, both key=1 + let outer_batch1 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![10])), + ], + )?; + let outer_batch2 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![2])), + Arc::new(Int32Array::from(vec![1])), // same key + Arc::new(Int32Array::from(vec![20])), + ], + )?; + + // Inner: 1 row, key=1 + let inner_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![100])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![50])), + ], + )?; + + let outer: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![outer_batch1, outer_batch2], + vec![false, true], // Pending before 2nd outer batch + )); + let inner: SendableRecordBatchStream = + Box::pin(PendingStream::new(vec![inner_batch], vec![false])); + + let on_outer: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + let on_inner: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + + let metrics = ExecutionPlanMetricsSet::new(); + let inner_schema = inner.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner_schema, &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + left_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + on_outer, + on_inner, + None, // no filter + LeftSemi, + 8192, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + )?; + + let batches = collect_stream(stream).await?; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 2, + "LeftSemi no filter: both outer rows (key=1) should be emitted \ + because inner has key=1. Got {total} rows." + ); + Ok(()) +} + +/// Verifies no-filter semi/anti joins when a matching outer key group spans +/// multiple batches and the next outer batch is temporarily unavailable. +/// +/// The outer input has an unmatched prefix row followed by a matching key +/// group that continues in the next batch. Both rows with key=1 should be +/// treated as matched. Returning `Pending` before the second batch makes the +/// join wait for the continuation while the key group is still open. +#[tokio::test] +async fn no_filter_boundary_pending_with_unmatched_prefix() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + + // Key=0 is unmatched. Key=1 matches inner and spans the batch boundary. + let outer_batch1 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![0, 1])), + Arc::new(Int32Array::from(vec![0, 1])), + Arc::new(Int32Array::from(vec![0, 10])), + ], + )?; + let outer_batch2 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![2])), + Arc::new(Int32Array::from(vec![1])), // same key + Arc::new(Int32Array::from(vec![20])), + ], + )?; + + // Key=1 matches two outer rows. Key=2 keeps the inner input non-exhausted. + let inner_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![100, 200])), + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![50, 60])), + ], + )?; + + let on_outer: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + let on_inner: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + + for (join_type, expected_a1) in [(LeftSemi, vec![1, 2]), (LeftAnti, vec![0])] { + let outer: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![outer_batch1.clone(), outer_batch2.clone()], + vec![false, true], // Pending before 2nd outer batch + )); + let inner: SendableRecordBatchStream = + Box::pin(PendingStream::new(vec![inner_batch.clone()], vec![false])); + + let metrics = ExecutionPlanMetricsSet::new(); + let inner_schema = inner.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner_schema, &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + Arc::clone(&left_schema), + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + on_outer.clone(), + on_inner.clone(), + None, // no filter + join_type, + 8192, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + )?; + + let batches = collect_stream(stream).await?; + let actual_a1 = batches + .iter() + .flat_map(|batch| { + let values = batch + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + (0..batch.num_rows()).map(|row| values.value(row)) + }) + .collect::<Vec<_>>(); + assert_eq!(actual_a1, expected_a1, "{join_type:?}"); + } + Ok(()) +} + +/// Same as the no-filter boundary case, with a filter: the outer key group +/// spans batches and the outer input returns Pending at the boundary. +/// +/// Setup: +/// - Outer: 2 single-row batches, both key=1, c1=[10, 20] +/// - Inner: 1 row, key=1, c2=10 +/// - Filter: c1 == c2 (first outer row matches, second doesn't) +/// - Pending before 2nd outer batch +/// +/// Expected: 1 row (only the first outer row c1=10 passes the filter) +#[tokio::test] +async fn filtered_boundary_pending_outer_rows() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + + let outer_batch1 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![10])), // matches filter + ], + )?; + let outer_batch2 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![2])), + Arc::new(Int32Array::from(vec![1])), // same key + Arc::new(Int32Array::from(vec![20])), // doesn't match + ], + )?; + + let inner_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![100])), + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![10])), + ], + )?; + + let outer: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![outer_batch1, outer_batch2], + vec![false, true], // Pending before 2nd outer batch + )); + let inner: SendableRecordBatchStream = + Box::pin(PendingStream::new(vec![inner_batch], vec![false])); + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Eq, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let on_outer: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + let on_inner: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + + let metrics = ExecutionPlanMetricsSet::new(); + let inner_schema = inner.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner_schema, &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + left_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + on_outer, + on_inner, + Some(filter), + LeftSemi, + 8192, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + )?; + + let batches = collect_stream(stream).await?; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total, 1, + "LeftSemi filtered boundary: only first outer row (c1=10) matches \ + filter c1==c2. Got {total} rows." + ); + Ok(()) +} + +// ── Bitwise stream spill tests ───────────────────────────────────────────── + +/// Exercises inner key group spilling under memory pressure. +/// +/// Uses a tiny memory limit (100 bytes) with disk spilling enabled. Since our +/// operator only buffers inner rows when a filter is present, this test includes +/// a filter (c1 < c2, always true). Verifies: +/// 1. Spill metrics are recorded (spill_count, spilled_bytes, spilled_rows > 0) +/// 2. Results match a non-spilled run +#[tokio::test] +async fn bitwise_spill_with_filter() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5, 6]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40, 50]), + ("b1", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + // c1 < c2 is always true for matching keys + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Lt, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + for batch_size in [1, 50] { + let session_config = SessionConfig::default().with_batch_size(batch_size); + + for join_type in [LeftSemi, LeftAnti, RightSemi, RightAnti] { + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config.clone()) + .with_runtime(Arc::clone(&runtime)), + ); + + let join = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + assert!( + join.metrics().is_some(), + "metrics missing for {join_type:?}" + ); + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}, batch_size={batch_size}" + ); + assert!( + metrics.spilled_bytes().unwrap() > 0, + "expected spilled_bytes > 0 for {join_type:?}, batch_size={batch_size}" + ); + assert!( + metrics.spilled_rows().unwrap() > 0, + "expected spilled_rows > 0 for {join_type:?}, batch_size={batch_size}" + ); + let join_time = metrics + .sum_by_name("join_time") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + join_time > 0, + "expected join_time > 0 for {join_type:?}, batch_size={batch_size}" + ); + let output_rows = metrics.output_rows().unwrap_or(0); + let collected_rows: usize = spilled_result.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + output_rows, collected_rows, + "output_rows metric should match collected rows for \ + {join_type:?}, batch_size={batch_size}" + ); + + // Run without spilling and compare results + let task_ctx_no_spill = Arc::new( + TaskContext::default().with_session_config(session_config.clone()), + ); + let join_no_spill = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + let no_spill_metrics = join_no_spill.metrics().unwrap(); + assert_eq!( + no_spill_metrics.spill_count(), + Some(0), + "unexpected spill for {join_type:?} without memory limit" + ); + + assert_eq!( + spilled_result, no_spill_result, + "spilled vs non-spilled results differ for {join_type:?}, batch_size={batch_size}" + ); + } + } + + Ok(()) +} + +/// A single inner key group spanning several inner batches can spill more +/// than once under memory pressure. Every spilled slice must still be +/// evaluated against the outer rows — an earlier spill file must not be +/// dropped when a later slice of the same group spills. +#[tokio::test] +async fn bitwise_multi_spill_inner_key_group() -> Result<()> { + // Outer: one row with key 1, c1 = 5. + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![5])); + + // Inner: one key group (b2 = 1) spanning two batches. Only the first + // batch satisfies the filter c1 < c2 (5 < 10); the second (5 < 0) does + // not, so dropping the first spilled slice flips the semi-join result. + let right_batches = vec![ + build_table_i32(("a2", &vec![10]), ("b2", &vec![1]), ("c2", &vec![10])), + build_table_i32(("a2", &vec![20]), ("b2", &vec![1]), ("c2", &vec![0])), + ]; + let right = build_table_from_batches(right_batches); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + let filter = build_c1_lt_c2_filter(left.schema().as_ref(), right.schema().as_ref()); + + // 100-byte pool: every buffered slice fails its reservation, so each + // inner batch of the key group spills separately. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::default().with_batch_size(1)) + .with_runtime(runtime), + ); + + let join = SortMergeJoinExec::try_new( + left, + right, + on, + Some(filter), + LeftSemi, + sort_options, + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let output_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + output_rows, 1, + "left row must match the group's first (spilled) inner slice", + ); + + let metrics = join.metrics().expect("must have metrics"); + assert_eq!( + metrics.spill_count(), + Some(1), + "all overflows of one key group must share a single spill file", + ); + assert_eq!( + metrics.spilled_rows(), + Some(2), + "both inner slices of the group must be spilled", + ); + Ok(()) +} + +/// Once the inner key group has spilled, an outer key group spanning a batch +/// boundary must still be evaluated against the spilled inner rows — the +/// second outer batch's rows must not be treated as having no inner group to +/// match against. +/// +/// Setup: +/// - Outer: 2 single-row batches, both key=1, c1=[10, 10] +/// - Inner: 1 batch with many rows all key=1 (enough to trigger spill) +/// - Filter: c1 == c2 (matches when c2=10) +/// - Memory limit: tiny (100 bytes) to force spilling +/// - Pending before 2nd outer batch, while the key group is still open +/// +/// Expected: both outer rows match (semi=2 rows, anti=0 rows) +#[tokio::test] +async fn spill_filtered_boundary_loses_outer_rows() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + + // Two single-row outer batches with the same key -- key group spans boundary + let outer_batch1 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(Int32Array::from(vec![1])), // key=1 + Arc::new(Int32Array::from(vec![10])), // matches filter + ], + )?; + let outer_batch2 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![2])), + Arc::new(Int32Array::from(vec![1])), // same key=1 + Arc::new(Int32Array::from(vec![10])), // also matches filter + ], + )?; + + // Inner: many rows with key=1 to force spilling, followed by key=2. + // c2=10 so the filter c1==c2 passes for both outer rows. + // The key=2 row ensures the inner cursor advances past the key group + // (buffer_inner_key_group returns Ok(false) instead of Ok(true)). + let n_inner = 200; + let mut inner_a = vec![100; n_inner]; + inner_a.push(101); + let mut inner_b = vec![1; n_inner]; + inner_b.push(2); // different key -- forces inner cursor past key=1 + let mut inner_c = vec![10; n_inner]; + inner_c.push(10); + let inner_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(inner_a)), + Arc::new(Int32Array::from(inner_b)), + Arc::new(Int32Array::from(inner_c)), + ], + )?; + + // Filter: c1 == c2 + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Eq, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let on_outer: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + let on_inner: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("b1", 1))]; + + for join_type in [LeftSemi, LeftAnti] { + let outer: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![outer_batch1.clone(), outer_batch2.clone()], + vec![false, true], // Pending before 2nd outer batch + )); + let inner: SendableRecordBatchStream = + Box::pin(PendingStream::new(vec![inner_batch.clone()], vec![false])); + + let metrics = ExecutionPlanMetricsSet::new(); + let reservation = MemoryConsumer::new("test").register(&runtime.memory_pool); + let spill_manager = SpillManager::new( + Arc::clone(&runtime), + SpillMetrics::new(&metrics, 0), + Arc::clone(&right_schema), + ); + + let stream = BitwiseSortMergeJoinStream::try_new( + Arc::clone(&left_schema), + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + on_outer.clone(), + on_inner.clone(), + Some(filter.clone()), + join_type, + 8192, + 0, + &metrics, + reservation, + spill_manager, + Arc::clone(&runtime), + )?; + + let batches = collect_stream(stream).await?; + let total: usize = batches.iter().map(|b| b.num_rows()).sum(); + + match join_type { + LeftSemi => { + assert_eq!( + total, 2, + "LeftSemi spill+boundary: both outer rows match filter, \ + expected 2 rows, got {total}" + ); + } + LeftAnti => { + assert_eq!( + total, 0, + "LeftAnti spill+boundary: both outer rows match filter, \ + expected 0 rows, got {total}" + ); + } + _ => unreachable!(), + } + } + + Ok(()) +} + +/// Verifies that `peak_mem_used` reflects spill read-back memory during +/// output materialization (multi-source path). +/// +/// When spilled buffered batches are read back from disk to produce join +/// output, a scoped `MemoryReservation` (via `new_empty()`) tracks the +/// transient memory. Its `Drop` guarantees the pool is balanced on every +/// exit path — normal return or early `?` error. +#[tokio::test] +async fn spill_read_back_memory_accounting() -> Result<()> { + use arrow::array::Array; + + let left_batch = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let size_estimation = left_batch.get_array_memory_size() + + Int32Array::from(vec![1, 1]).get_array_memory_size() + + 2usize.next_power_of_two() * size_of::<usize>() + + size_of::<std::ops::Range<usize>>() + + size_of::<usize>(); + + // Memory limit too small for a full batch — forces spilling. + let memory_limit = size_estimation / 2; + + // All rows share the same join key (b=1) to force multiple buffered + // batches in the same key group — triggering spill read-back during + // output materialization. + let left_batches: Vec<RecordBatch> = (0..4) + .map(|i| { + build_table_i32( + ("a1", &vec![i * 2, i * 2 + 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![100 + i, 101 + i]), + ) + }) + .collect(); + let left = build_table_from_batches(left_batches); + + let right_batches: Vec<RecordBatch> = (0..4) + .map(|i| { + build_table_i32( + ("a2", &vec![i * 2, i * 2 + 1]), + ("b2", &vec![1, 1]), + ("c2", &vec![200 + i, 201 + i]), + ) + }) + .collect(); + let right = build_table_from_batches(right_batches); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let session_config = SessionConfig::default().with_batch_size(50); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime)), + ); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Inner, + sort_options, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let result = common::collect(stream).await.unwrap(); + + assert!(!result.is_empty(), "Expected non-empty join result"); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur" + ); + + // peak_mem_used should reflect the spill read-back: when buffered + // batches are read from disk during output materialization, grow() + // temporarily reserves size_estimation. This pushes peak above what + // join_arrays_mem alone would show. + let peak_mem = metrics + .sum_by_name("peak_mem_used") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + peak_mem >= size_estimation, + "peak_mem_used ({peak_mem}) should be >= size_estimation ({size_estimation}) \ + because spill read-back temporarily loads full batch into memory" + ); + + // All memory must be released (grow/shrink balanced) + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "All memory should be released after join completes" + ); + + Ok(()) +} + +/// Verifies spill read-back memory tracking for the single-source path. +/// +/// When only ONE buffered batch exists for a key group and it's spilled, +/// `fetch_right_columns_by_idxs` reads it back. A scoped `MemoryReservation` +/// (via `new_empty()`) tracks the transient memory and releases it on drop. +#[tokio::test] +async fn spill_read_back_single_source() -> Result<()> { + use arrow::array::Array; + + let left_batch = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let size_estimation = left_batch.get_array_memory_size() + + Int32Array::from(vec![1, 1]).get_array_memory_size() + + 2usize.next_power_of_two() * size_of::<usize>() + + size_of::<std::ops::Range<usize>>() + + size_of::<usize>(); + + // Memory limit too small for a full batch — forces spilling. + let memory_limit = size_estimation / 2; + + // Multiple distinct keys so each key group has exactly ONE buffered batch. + // This ensures the single-source path is exercised. + let left_batches: Vec<RecordBatch> = (0..4) + .map(|i| { + build_table_i32( + ("a1", &vec![i * 2, i * 2 + 1]), + ("b1", &vec![i, i]), + ("c1", &vec![100 + i, 101 + i]), + ) + }) + .collect(); + let left = build_table_from_batches(left_batches); + + // One batch per key — each key group has single source + let right_batches: Vec<RecordBatch> = (0..4) + .map(|i| { + build_table_i32( + ("a2", &vec![i * 2, i * 2 + 1]), + ("b2", &vec![i, i]), + ("c2", &vec![200 + i, 201 + i]), + ) + }) + .collect(); + let right = build_table_from_batches(right_batches); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + + let session_config = SessionConfig::default().with_batch_size(50); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(Arc::clone(&runtime)), + ); + + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Inner, + sort_options, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let result = common::collect(stream).await.unwrap(); + + assert!(!result.is_empty(), "Expected non-empty join result"); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur" + ); + + // peak_mem_used should reflect the single-batch read-back + let peak_mem = metrics + .sum_by_name("peak_mem_used") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + peak_mem >= size_estimation, + "peak_mem_used ({peak_mem}) should be >= size_estimation ({size_estimation}) \ + because single-source spill read-back loads full batch" + ); + + // All memory must be released + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "All memory should be released after join completes" + ); + + Ok(()) +} + +/// Small chunk size so even tiny test spill files are split into several +/// pieces, forcing multiple genuine suspend/resume cycles instead of one. +const PENDING_CHUNK_SIZE: usize = 16; + +/// Splits real spill bytes into fixed-size chunks and yields `Poll::Pending` +/// before every chunk +struct PendingChunkedStream { + chunks: VecDeque<Bytes>, + yield_pending: bool, +} + +impl PendingChunkedStream { + fn new(bytes: Bytes) -> Self { + let mut chunks = VecDeque::new(); + if bytes.is_empty() { + chunks.push_back(bytes); + } else { + let mut remaining = bytes; + while !remaining.is_empty() { + let take = PENDING_CHUNK_SIZE.min(remaining.len()); + chunks.push_back(remaining.split_to(take)); + } + } + Self { + chunks, + yield_pending: true, + } + } +} + +impl Stream for PendingChunkedStream { + type Item = Result<Bytes>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.yield_pending { + self.yield_pending = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + // Pending before every subsequent chunk as well. + self.yield_pending = true; + match self.chunks.pop_front() { + Some(chunk) => Poll::Ready(Some(Ok(chunk))), + None => Poll::Ready(None), + } + } +} + +/// A `SpillFile` that delegates everything to a real local spill file, +/// except `read_stream`, which is forced through `PendingChunkedStream`. +struct PendingSpillFile { + inner: Arc<dyn SpillFile>, +} + +impl SpillFile for PendingSpillFile { + fn path(&self) -> Option<&std::path::Path> { + self.inner.path() + } + + fn size(&self) -> Option<u64> { + self.inner.size() + } + + fn read_stream(&self) -> Result<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>> { + let path = self + .inner + .path() + .expect("PendingSpillFile only wraps local files") + .to_owned(); + + let stream = futures::stream::once(async move { + tokio::fs::read(&path) + .await + .map(Bytes::from) + .map_err(datafusion_common::DataFusionError::IoError) + }) + .flat_map( + |read_result| -> Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>> { + match read_result { + Ok(bytes) => Box::pin(PendingChunkedStream::new(bytes)), + Err(e) => Box::pin(futures::stream::once(async move { Err(e) })), + } + }, + ); + + Ok(Box::pin(stream)) + } + + fn open_writer(&self) -> Result<Box<dyn SpillWriter>> { + self.inner.open_writer() + } +} + +/// Wraps the default `OsTmpDirectory` factory so every spill file it +/// creates is a [`PendingSpillFile`]. +struct PendingTempFileFactory { + inner: Arc<DiskManager>, +} + +impl TempFileFactory for PendingTempFileFactory { + fn create_temp_file(&self, description: &str) -> Result<Arc<dyn SpillFile>> { + Ok(Arc::new(PendingSpillFile { + inner: self.inner.create_tmp_file(description)?, + })) + } +} + +fn pending_disk_manager_builder() -> DiskManagerBuilder { + let inner = Arc::new( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .build() + .unwrap(), + ); + DiskManagerBuilder::default().with_mode(DiskManagerMode::Custom(Arc::new( + PendingTempFileFactory { inner }, + ))) +} + +/// Materializing-side (Inner/Left/Right/Full) coverage: identical to +/// `overallocation_multi_batch_spill`, but every spill read goes through +/// `PendingSpillFile`, so `poll_spilled_batches` must actually hit and +/// recover from `Poll::Pending` mid-read. +#[tokio::test] +async fn materializing_spill_pending_stream() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![1, 1]), + ("c1", &vec![6, 7]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10]), + ("b2", &vec![1, 1]), + ("c2", &vec![50, 60]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![20, 30]), + ("b2", &vec![1, 1]), + ("c2", &vec![70, 80]), + ); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2]); + let right = build_table_from_batches(vec![right_batch_1, right_batch_2]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(500, 1.0) + .with_disk_manager_builder(pending_disk_manager_builder()) + .build_arc()?; + + for join_type in [Inner, Left, Right, Full] { + let task_ctx = + Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}" + ); + + // Compare against a no-spill run to make sure waiting on the + // spill reads didn't corrupt or drop any data. + let task_ctx_no_spill = Arc::new(TaskContext::default()); + let join_no_spill = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + assert_eq!( + spilled_result, no_spill_result, + "Pending-forced spill read produced different results for {join_type:?}" + ); + } + + Ok(()) +} + +/// Bitwise-side (Semi/Anti) coverage: identical to `bitwise_spill_with_filter`, +/// but every spill read goes through `PendingSpillFile`, so reading the +/// spilled inner rows back must actually hit and recover from `Poll::Pending` +/// mid-read. +#[tokio::test] +async fn bitwise_spill_pending_stream() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5, 6]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40, 50]), + ("b1", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + // c1 < c2 is always true for matching keys — same filter as + // bitwise_spill_with_filter, so the inner key group is buffered + // (and spilled) rather than short-circuited. + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Lt, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder(pending_disk_manager_builder()) + .build_arc()?; + + for join_type in [LeftSemi, LeftAnti, RightSemi, RightAnti] { + let task_ctx = + Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let join = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}" + ); + + let task_ctx_no_spill = Arc::new(TaskContext::default()); + let join_no_spill = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + assert_eq!( + spilled_result, no_spill_result, + "Pending-forced spill read produced different results for {join_type:?}" + ); + } + + Ok(()) +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/stream_join_utils.rs b/native/vendor/datafusion-physical-plan/src/joins/stream_join_utils.rs new file mode 100644 index 0000000..05a56d2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/stream_join_utils.rs
@@ -0,0 +1,1185 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This file contains common subroutines for symmetric hash join +//! related functionality, used both in join calculations and optimization rules. + +use std::collections::{HashMap, VecDeque}; +use std::mem::size_of; +use std::sync::Arc; + +use crate::joins::MapOffset; +use crate::joins::join_hash_map::{ + contain_hashes, get_matched_indices, get_matched_indices_with_limit_offset, + update_from_iter, +}; +use crate::joins::utils::{JoinFilter, JoinHashMapType}; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, +}; +use crate::{ExecutionPlan, metrics}; + +use arrow::array::{ + ArrowPrimitiveType, BooleanArray, BooleanBufferBuilder, NativeAdapter, + PrimitiveArray, RecordBatch, +}; +use arrow::buffer::NullBuffer; +use arrow::compute::concat_batches; +use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef}; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::utils::memory::estimate_memory_size; +use datafusion_common::{HashSet, JoinSide, Result, ScalarValue, arrow_datafusion_err}; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::intervals::cp_solver::ExprIntervalGraph; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr}; + +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use hashbrown::HashTable; + +/// Implementation of `JoinHashMapType` for `PruningJoinHashMap`. +impl JoinHashMapType for PruningJoinHashMap { + // Extend with zero + fn extend_zero(&mut self, len: usize) { + self.next.resize(self.next.len() + len, 0) + } + + fn update_from_iter<'a>( + &mut self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + Send + 'a>, + deleted_offset: usize, + ) { + let slice: &mut [u64] = self.next.make_contiguous(); + update_from_iter::<u64>(&mut self.map, slice, iter, deleted_offset); + } + + fn get_matched_indices<'a>( + &self, + iter: Box<dyn Iterator<Item = (usize, &'a u64)> + 'a>, + deleted_offset: Option<usize>, + ) -> (Vec<u32>, Vec<u64>) { + // Flatten the deque + let next: Vec<u64> = self.next.iter().copied().collect(); + get_matched_indices::<u64>(&self.map, &next, iter, deleted_offset) + } + + fn get_matched_indices_with_limit_offset( + &self, + hash_values: &[u64], + valid_keys: Option<&NullBuffer>, + limit: usize, + offset: MapOffset, + input_indices: &mut Vec<u32>, + match_indices: &mut Vec<u64>, + ) -> Option<MapOffset> { + // Flatten the deque + let next: Vec<u64> = self.next.iter().copied().collect(); + get_matched_indices_with_limit_offset::<u64>( + &self.map, + &next, + hash_values, + valid_keys, + limit, + offset, + input_indices, + match_indices, + ) + } + + fn contain_hashes(&self, hash_values: &[u64]) -> BooleanArray { + contain_hashes(&self.map, hash_values) + } + + fn is_empty(&self) -> bool { + self.map.is_empty() + } + + fn len(&self) -> usize { + self.map.len() + } +} + +/// The `PruningJoinHashMap` is similar to a regular `JoinHashMap`, but with +/// the capability of pruning elements in an efficient manner. This structure +/// is particularly useful for cases where it's necessary to remove elements +/// from the map based on their buffer order. +/// +/// # Example +/// +/// ``` text +/// Let's continue the example of `JoinHashMap` and then show how `PruningJoinHashMap` would +/// handle the pruning scenario. +/// +/// Insert the pair (10,4) into the `PruningJoinHashMap`: +/// map: +/// ---------- +/// | 10 | 5 | +/// | 20 | 3 | +/// ---------- +/// list: +/// --------------------- +/// | 0 | 0 | 0 | 2 | 4 | <--- hash value 10 maps to 5,4,2 (which means indices values 4,3,1) +/// --------------------- +/// +/// Now, let's prune 3 rows from `PruningJoinHashMap`: +/// map: +/// --------- +/// | 1 | 5 | +/// --------- +/// list: +/// --------- +/// | 2 | 4 | <--- hash value 10 maps to 2 (5 - 3), 1 (4 - 3), NA (2 - 3) (which means indices values 1,0) +/// --------- +/// +/// After pruning, the | 2 | 3 | entry is deleted from `PruningJoinHashMap` since +/// there are no values left for this key. +/// ``` +pub struct PruningJoinHashMap { + /// Stores hash value to last row index + pub map: HashTable<(u64, u64)>, + /// Stores indices in chained list data structure + pub next: VecDeque<u64>, +} + +impl PruningJoinHashMap { + /// Constructs a new `PruningJoinHashMap` with the given capacity. + /// Both the map and the list are pre-allocated with the provided capacity. + /// + /// # Arguments + /// * `capacity`: The initial capacity of the hash map. + /// + /// # Returns + /// A new instance of `PruningJoinHashMap`. + pub(crate) fn with_capacity(capacity: usize) -> Self { + PruningJoinHashMap { + map: HashTable::with_capacity(capacity), + next: VecDeque::with_capacity(capacity), + } + } + + /// Shrinks the capacity of the hash map, if necessary, based on the + /// provided scale factor. + /// + /// # Arguments + /// * `scale_factor`: The scale factor that determines how conservative the + /// shrinking strategy is. The capacity will be reduced by 1/`scale_factor` + /// when necessary. + /// + /// # Note + /// Increasing the scale factor results in less aggressive capacity shrinking, + /// leading to potentially higher memory usage but fewer resizes. Conversely, + /// decreasing the scale factor results in more aggressive capacity shrinking, + /// potentially leading to lower memory usage but more frequent resizing. + pub(crate) fn shrink_if_necessary(&mut self, scale_factor: usize) { + let capacity = self.map.capacity(); + + if capacity > scale_factor * self.map.len() { + let new_capacity = (capacity * (scale_factor - 1)) / scale_factor; + // Resize the map with the new capacity. + self.map.shrink_to(new_capacity, |(hash, _)| *hash) + } + } + + /// Calculates the size of the `PruningJoinHashMap` in bytes. + /// + /// # Returns + /// The size of the hash map in bytes. + pub(crate) fn size(&self) -> usize { + let fixed_size = size_of::<PruningJoinHashMap>(); + + // TODO: switch to using [HashTable::allocation_size] when available after upgrading hashbrown to 0.15 + estimate_memory_size::<(u64, u64)>(self.map.capacity(), fixed_size).unwrap() + + self.next.capacity() * size_of::<u64>() + } + + /// Removes hash values from the map and the list based on the given pruning + /// length and deleting offset. + /// + /// # Arguments + /// * `prune_length`: The number of elements to remove from the list. + /// * `deleting_offset`: The offset used to determine which hash values to remove from the map. + /// + /// # Returns + /// A `Result` indicating whether the operation was successful. + pub(crate) fn prune_hash_values( + &mut self, + prune_length: usize, + deleting_offset: u64, + shrink_factor: usize, + ) { + // Remove elements from the list based on the pruning length. + self.next.drain(0..prune_length); + + // Calculate the keys that should be removed from the map. + let removable_keys = self + .map + .iter() + .filter_map(|(hash, tail_index)| { + (*tail_index < prune_length as u64 + deleting_offset).then_some(*hash) + }) + .collect::<Vec<_>>(); + + // Remove the keys from the map. + removable_keys.into_iter().for_each(|hash_value| { + self.map + .find_entry(hash_value, |(hash, _)| hash_value == *hash) + .unwrap() + .remove(); + }); + + // Shrink the map if necessary. + self.shrink_if_necessary(shrink_factor); + } +} + +fn check_filter_expr_contains_sort_information( + expr: &Arc<dyn PhysicalExpr>, + reference: &Arc<dyn PhysicalExpr>, +) -> bool { + expr.eq(reference) + || expr + .children() + .iter() + .any(|e| check_filter_expr_contains_sort_information(e, reference)) +} + +/// Create a one to one mapping from main columns to filter columns using +/// filter column indices. A column index looks like: +/// ```text +/// ColumnIndex { +/// index: 0, // field index in main schema +/// side: JoinSide::Left, // child side +/// } +/// ``` +pub fn map_origin_col_to_filter_col( + filter: &JoinFilter, + schema: &SchemaRef, + side: &JoinSide, +) -> Result<HashMap<Column, Column>> { + let filter_schema = filter.schema(); + let mut col_to_col_map = HashMap::<Column, Column>::new(); + for (filter_schema_index, index) in filter.column_indices().iter().enumerate() { + if index.side.eq(side) { + // Get the main field from column index: + let main_field = schema.field(index.index); + // Create a column expression: + let main_col = Column::new_with_schema(main_field.name(), schema.as_ref())?; + // Since the order of by filter.column_indices() is the same with + // that of intermediate schema fields, we can get the column directly. + let filter_field = filter_schema.field(filter_schema_index); + let filter_col = Column::new(filter_field.name(), filter_schema_index); + // Insert mapping: + col_to_col_map.insert(main_col, filter_col); + } + } + Ok(col_to_col_map) +} + +/// This function analyzes [`PhysicalSortExpr`] graphs with respect to output orderings +/// (sorting) properties. This is necessary since monotonically increasing and/or +/// decreasing expressions are required when using join filter expressions for +/// data pruning purposes. +/// +/// The method works as follows: +/// 1. Maps the original columns to the filter columns using the [`map_origin_col_to_filter_col`] function. +/// 2. Collects all columns in the sort expression using the [`collect_columns`] function. +/// 3. Checks if all columns are included in the map we obtain in the first step. +/// 4. If all columns are included, the sort expression is converted into a filter expression using +/// the [`convert_filter_columns`] function. +/// 5. Searches for the converted filter expression in the filter expression using the +/// [`check_filter_expr_contains_sort_information`] function. +/// 6. If an exact match is found, returns the converted filter expression as `Some(Arc<dyn PhysicalExpr>)`. +/// 7. If all columns are not included or an exact match is not found, returns [`None`]. +/// +/// Examples: +/// Consider the filter expression "a + b > c + 10 AND a + b < c + 100". +/// 1. If the expression "a@ + d@" is sorted, it will not be accepted since the "d@" column is not part of the filter. +/// 2. If the expression "d@" is sorted, it will not be accepted since the "d@" column is not part of the filter. +/// 3. If the expression "a@ + b@ + c@" is sorted, all columns are represented in the filter expression. However, +/// there is no exact match, so this expression does not indicate pruning. +pub fn convert_sort_expr_with_filter_schema( + side: &JoinSide, + filter: &JoinFilter, + schema: &SchemaRef, + sort_expr: &PhysicalSortExpr, +) -> Result<Option<Arc<dyn PhysicalExpr>>> { + let column_map = map_origin_col_to_filter_col(filter, schema, side)?; + let expr = Arc::clone(&sort_expr.expr); + // Get main schema columns: + let expr_columns = collect_columns(&expr); + // Calculation is possible with `column_map` since sort exprs belong to a child. + let all_columns_are_included = + expr_columns.iter().all(|col| column_map.contains_key(col)); + if all_columns_are_included { + // Since we are sure that one to one column mapping includes all columns, we convert + // the sort expression into a filter expression. + let converted_filter_expr = expr + .transform_up(|p| { + convert_filter_columns(p.as_ref(), &column_map).map(|transformed| { + match transformed { + Some(transformed) => Transformed::yes(transformed), + None => Transformed::no(p), + } + }) + }) + .data()?; + // Search the converted `PhysicalExpr` in filter expression; if an exact + // match is found, use this sorted expression in graph traversals. + if check_filter_expr_contains_sort_information( + filter.expression(), + &converted_filter_expr, + ) { + return Ok(Some(converted_filter_expr)); + } + } + Ok(None) +} + +/// This function is used to build the filter expression based on the sort order of input columns. +/// +/// It first calls the [`convert_sort_expr_with_filter_schema`] method to determine if the sort +/// order of columns can be used in the filter expression. If it returns a [`Some`] value, the +/// method wraps the result in a [`SortedFilterExpr`] instance with the original sort expression and +/// the converted filter expression. Otherwise, this function returns an error. +/// +/// The `SortedFilterExpr` instance contains information about the sort order of columns that can +/// be used in the filter expression, which can be used to optimize the query execution process. +pub fn build_filter_input_order( + side: JoinSide, + filter: &JoinFilter, + schema: &SchemaRef, + order: &PhysicalSortExpr, +) -> Result<Option<SortedFilterExpr>> { + let opt_expr = convert_sort_expr_with_filter_schema(&side, filter, schema, order)?; + opt_expr + .map(|filter_expr| { + SortedFilterExpr::try_new(order.clone(), filter_expr, filter.schema()) + }) + .transpose() +} + +/// Convert a physical expression into a filter expression using the given +/// column mapping information. +fn convert_filter_columns( + input: &dyn PhysicalExpr, + column_map: &HashMap<Column, Column>, +) -> Result<Option<Arc<dyn PhysicalExpr>>> { + // Attempt to downcast the input expression to a Column type. + Ok(if let Some(col) = input.downcast_ref::<Column>() { + // If the downcast is successful, retrieve the corresponding filter column. + column_map.get(col).map(|c| Arc::new(c.clone()) as _) + } else { + // If the downcast fails, return the input expression as is. + None + }) +} + +/// The [SortedFilterExpr] object represents a sorted filter expression. It +/// contains the following information: The origin expression, the filter +/// expression, an interval encapsulating expression bounds, and a stable +/// index identifying the expression in the expression DAG. +/// +/// Physical schema of a [JoinFilter]'s intermediate batch combines two sides +/// and uses new column names. In this process, a column exchange is done so +/// we can utilize sorting information while traversing the filter expression +/// DAG for interval calculations. When evaluating the inner buffer, we use +/// `origin_sorted_expr`. +#[derive(Debug, Clone)] +pub struct SortedFilterExpr { + /// Sorted expression from a join side (i.e. a child of the join) + origin_sorted_expr: PhysicalSortExpr, + /// Expression adjusted for filter schema. + filter_expr: Arc<dyn PhysicalExpr>, + /// Interval containing expression bounds + interval: Interval, + /// Node index in the expression DAG + node_index: usize, +} + +impl SortedFilterExpr { + /// Constructor + pub fn try_new( + origin_sorted_expr: PhysicalSortExpr, + filter_expr: Arc<dyn PhysicalExpr>, + filter_schema: &Schema, + ) -> Result<Self> { + let dt = filter_expr.data_type(filter_schema)?; + Ok(Self { + origin_sorted_expr, + filter_expr, + interval: Interval::make_unbounded(&dt)?, + node_index: 0, + }) + } + + /// Get origin expr information + pub fn origin_sorted_expr(&self) -> &PhysicalSortExpr { + &self.origin_sorted_expr + } + + /// Get filter expr information + pub fn filter_expr(&self) -> &Arc<dyn PhysicalExpr> { + &self.filter_expr + } + + /// Get interval information + pub fn interval(&self) -> &Interval { + &self.interval + } + + /// Sets interval + pub fn set_interval(&mut self, interval: Interval) { + self.interval = interval; + } + + /// Node index in ExprIntervalGraph + pub fn node_index(&self) -> usize { + self.node_index + } + + /// Node index setter in ExprIntervalGraph + pub fn set_node_index(&mut self, node_index: usize) { + self.node_index = node_index; + } +} + +/// Calculate the filter expression intervals. +/// +/// This function updates the `interval` field of each `SortedFilterExpr` based +/// on the first or the last value of the expression in `build_input_buffer` +/// and `probe_batch`. +/// +/// # Parameters +/// +/// * `build_input_buffer` - The [RecordBatch] on the build side of the join. +/// * `build_sorted_filter_expr` - Build side [SortedFilterExpr] to update. +/// * `probe_batch` - The `RecordBatch` on the probe side of the join. +/// * `probe_sorted_filter_expr` - Probe side `SortedFilterExpr` to update. +/// +/// ## Note +/// +/// Utilizing interval arithmetic, this function computes feasible join intervals +/// on the pruning side by evaluating the prospective value ranges that might +/// emerge in subsequent data batches from the enforcer side. This is done by +/// first creating an interval for join filter values in the pruning side of the +/// join, which spans `[-∞, FV]` or `[FV, ∞]` depending on the ordering (descending/ +/// ascending) of the filter expression. Here, `FV` denotes the first value on the +/// pruning side. This range is then compared with the enforcer side interval, +/// which either spans `[-∞, LV]` or `[LV, ∞]` depending on the ordering (ascending/ +/// descending) of the probe side. Here, `LV` denotes the last value on the enforcer +/// side. +/// +/// As a concrete example, consider the following query: +/// +/// ```text +/// SELECT * FROM left_table, right_table +/// WHERE +/// left_key = right_key AND +/// a > b - 3 AND +/// a < b + 10 +/// ``` +/// +/// where columns `a` and `b` come from tables `left_table` and `right_table`, +/// respectively. When a new `RecordBatch` arrives at the right side, the +/// condition `a > b - 3` will possibly indicate a prunable range for the left +/// side. Conversely, when a new `RecordBatch` arrives at the left side, the +/// condition `a < b + 10` will possibly indicate prunability for the right side. +/// Let’s inspect what happens when a new `RecordBatch` arrives at the right +/// side (i.e. when the left side is the build side): +/// +/// ```text +/// Build Probe +/// +-------+ +-------+ +/// | a | z | | b | y | +/// |+--|--+| |+--|--+| +/// | 1 | 2 | | 4 | 3 | +/// |+--|--+| |+--|--+| +/// | 3 | 1 | | 4 | 3 | +/// |+--|--+| |+--|--+| +/// | 5 | 7 | | 6 | 1 | +/// |+--|--+| |+--|--+| +/// | 7 | 1 | | 6 | 3 | +/// +-------+ +-------+ +/// ``` +/// +/// In this case, the interval representing viable (i.e. joinable) values for +/// column `a` is `[1, ∞]`, and the interval representing possible future values +/// for column `b` is `[6, ∞]`. With these intervals at hand, we next calculate +/// intervals for the whole filter expression and propagate join constraint by +/// traversing the expression graph. +pub fn calculate_filter_expr_intervals( + build_input_buffer: &RecordBatch, + build_sorted_filter_expr: &mut SortedFilterExpr, + probe_batch: &RecordBatch, + probe_sorted_filter_expr: &mut SortedFilterExpr, +) -> Result<()> { + // If either build or probe side has no data, return early: + if build_input_buffer.num_rows() == 0 || probe_batch.num_rows() == 0 { + return Ok(()); + } + // Calculate the interval for the build side filter expression (if present): + update_filter_expr_interval( + &build_input_buffer.slice(0, 1), + build_sorted_filter_expr, + )?; + // Calculate the interval for the probe side filter expression (if present): + update_filter_expr_interval( + &probe_batch.slice(probe_batch.num_rows() - 1, 1), + probe_sorted_filter_expr, + ) +} + +/// This is a subroutine of the function [`calculate_filter_expr_intervals`]. +/// It constructs the current interval using the given `batch` and updates +/// the filter expression (i.e. `sorted_expr`) with this interval. +pub fn update_filter_expr_interval( + batch: &RecordBatch, + sorted_expr: &mut SortedFilterExpr, +) -> Result<()> { + // Evaluate the filter expression and convert the result to an array: + let array = sorted_expr + .origin_sorted_expr() + .expr + .evaluate(batch)? + .into_array(1)?; + // Convert the array to a ScalarValue: + let value = ScalarValue::try_from_array(&array, 0)?; + // Create a ScalarValue representing positive or negative infinity for the same data type: + let inf = ScalarValue::try_from(value.data_type())?; + // Update the interval with lower and upper bounds based on the sort option: + let interval = if sorted_expr.origin_sorted_expr().options.descending { + Interval::try_new(inf, value)? + } else { + Interval::try_new(value, inf)? + }; + // Set the calculated interval for the sorted filter expression: + sorted_expr.set_interval(interval); + Ok(()) +} + +/// Get the anti join indices from the visited hash set. +/// +/// This method returns the indices from the original input that were not present in the visited hash set. +/// +/// # Arguments +/// +/// * `prune_length` - The length of the pruned record batch. +/// * `deleted_offset` - The offset to the indices. +/// * `visited_rows` - The hash set of visited indices. +/// +/// # Returns +/// +/// A `PrimitiveArray` of the anti join indices. +pub fn get_pruning_anti_indices<T: ArrowPrimitiveType>( + prune_length: usize, + deleted_offset: usize, + visited_rows: &HashSet<usize>, +) -> PrimitiveArray<T> +where + NativeAdapter<T>: From<<T as ArrowPrimitiveType>::Native>, +{ + let mut bitmap = BooleanBufferBuilder::new(prune_length); + bitmap.append_n(prune_length, false); + // mark the indices as true if they are present in the visited hash set + for v in 0..prune_length { + let row = v + deleted_offset; + bitmap.set_bit(v, visited_rows.contains(&row)); + } + // get the anti index + (0..prune_length) + .filter_map(|idx| (!bitmap.get_bit(idx)).then_some(T::Native::from_usize(idx))) + .collect() +} + +/// This method creates a boolean buffer from the visited rows hash set +/// and the indices of the pruned record batch slice. +/// +/// It gets the indices from the original input that were present in the visited hash set. +/// +/// # Arguments +/// +/// * `prune_length` - The length of the pruned record batch. +/// * `deleted_offset` - The offset to the indices. +/// * `visited_rows` - The hash set of visited indices. +/// +/// # Returns +/// +/// A [PrimitiveArray] of the specified type T, containing the semi indices. +pub fn get_pruning_semi_indices<T: ArrowPrimitiveType>( + prune_length: usize, + deleted_offset: usize, + visited_rows: &HashSet<usize>, +) -> PrimitiveArray<T> +where + NativeAdapter<T>: From<<T as ArrowPrimitiveType>::Native>, +{ + let mut bitmap = BooleanBufferBuilder::new(prune_length); + bitmap.append_n(prune_length, false); + // mark the indices as true if they are present in the visited hash set + (0..prune_length).for_each(|v| { + let row = &(v + deleted_offset); + bitmap.set_bit(v, visited_rows.contains(row)); + }); + // get the semi index + (0..prune_length) + .filter_map(|idx| (bitmap.get_bit(idx)).then_some(T::Native::from_usize(idx))) + .collect() +} + +pub fn combine_two_batches( + output_schema: &SchemaRef, + left_batch: Option<RecordBatch>, + right_batch: Option<RecordBatch>, +) -> Result<Option<RecordBatch>> { + match (left_batch, right_batch) { + (Some(batch), None) | (None, Some(batch)) => { + // If only one of the batches are present, return it: + Ok(Some(batch)) + } + (Some(left_batch), Some(right_batch)) => { + // If both batches are present, concatenate them: + concat_batches(output_schema, &[left_batch, right_batch]) + .map_err(|e| arrow_datafusion_err!(e)) + .map(Some) + } + (None, None) => { + // If neither is present, return an empty batch: + Ok(None) + } + } +} + +/// Records the visited indices from the input `PrimitiveArray` of type `T` into the given hash set `visited`. +/// This function will insert the indices (offset by `offset`) into the `visited` hash set. +/// +/// # Arguments +/// +/// * `visited` - A hash set to store the visited indices. +/// * `offset` - An offset to the indices in the `PrimitiveArray`. +/// * `indices` - The input `PrimitiveArray` of type `T` which stores the indices to be recorded. +pub fn record_visited_indices<T: ArrowPrimitiveType>( + visited: &mut HashSet<usize>, + offset: usize, + indices: &PrimitiveArray<T>, +) { + for i in indices.values() { + visited.insert(i.as_usize() + offset); + } +} + +#[derive(Debug)] +pub struct StreamJoinSideMetrics { + /// Number of batches consumed by this operator + pub(crate) input_batches: metrics::Count, + /// Number of rows consumed by this operator + pub(crate) input_rows: metrics::Count, +} + +/// Metrics for HashJoinExec +#[derive(Debug)] +pub struct StreamJoinMetrics { + /// Number of left batches/rows consumed by this operator + pub(crate) left: StreamJoinSideMetrics, + /// Number of right batches/rows consumed by this operator + pub(crate) right: StreamJoinSideMetrics, + /// Memory used by sides in bytes + pub(crate) stream_memory_usage: metrics::Gauge, + /// Number of rows produced by this operator + pub(crate) baseline_metrics: BaselineMetrics, +} + +impl StreamJoinMetrics { + pub fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + let input_batches = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("left_input_batches", partition); + let input_rows = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("left_input_rows", partition); + let left = StreamJoinSideMetrics { + input_batches, + input_rows, + }; + + let input_batches = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("right_input_batches", partition); + let input_rows = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("right_input_rows", partition); + let right = StreamJoinSideMetrics { + input_batches, + input_rows, + }; + + let stream_memory_usage = MetricBuilder::new(metrics) + .with_category(MetricCategory::Bytes) + .gauge("stream_memory_usage", partition); + + Self { + left, + right, + stream_memory_usage, + baseline_metrics: BaselineMetrics::new(metrics, partition), + } + } +} + +/// Updates sorted filter expressions with corresponding node indices from the +/// expression interval graph. +/// +/// This function iterates through the provided sorted filter expressions, +/// gathers the corresponding node indices from the expression interval graph, +/// and then updates the sorted expressions with these indices. It ensures +/// that these sorted expressions are aligned with the structure of the graph. +fn update_sorted_exprs_with_node_indices( + graph: &mut ExprIntervalGraph, + sorted_exprs: &mut [SortedFilterExpr], +) { + // Extract filter expressions from the sorted expressions: + let filter_exprs = sorted_exprs + .iter() + .map(|expr| Arc::clone(expr.filter_expr())) + .collect::<Vec<_>>(); + + // Gather corresponding node indices for the extracted filter expressions from the graph: + let child_node_indices = graph.gather_node_indices(&filter_exprs); + + // Iterate through the sorted expressions and the gathered node indices: + for (sorted_expr, (_, index)) in sorted_exprs.iter_mut().zip(child_node_indices) { + // Update each sorted expression with the corresponding node index: + sorted_expr.set_node_index(index); + } +} + +/// Prepares and sorts expressions based on a given filter, left and right schemas, +/// and sort expressions. +/// +/// This function prepares sorted filter expressions for both the left and right +/// sides of a join operation. It first builds the filter order for each side +/// based on the provided `ExecutionPlan`. If both sides have valid sorted filter +/// expressions, the function then constructs an expression interval graph and +/// updates the sorted expressions with node indices. The final sorted filter +/// expressions for both sides are then returned. +/// +/// # Parameters +/// +/// * `filter` - The join filter to base the sorting on. +/// * `left` - The `ExecutionPlan` for the left side of the join. +/// * `right` - The `ExecutionPlan` for the right side of the join. +/// * `left_sort_exprs` - The expressions to sort on the left side. +/// * `right_sort_exprs` - The expressions to sort on the right side. +/// +/// # Returns +/// +/// * A tuple consisting of the sorted filter expression for the left and right sides, and an expression interval graph. +pub fn prepare_sorted_exprs( + filter: &JoinFilter, + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + left_sort_exprs: &LexOrdering, + right_sort_exprs: &LexOrdering, +) -> Result<(SortedFilterExpr, SortedFilterExpr, ExprIntervalGraph)> { + let err = || { + datafusion_common::plan_datafusion_err!("Filter does not include the child order") + }; + + // Build the filter order for the left side: + let left_temp_sorted_filter_expr = build_filter_input_order( + JoinSide::Left, + filter, + &left.schema(), + &left_sort_exprs[0], + )? + .ok_or_else(err)?; + + // Build the filter order for the right side: + let right_temp_sorted_filter_expr = build_filter_input_order( + JoinSide::Right, + filter, + &right.schema(), + &right_sort_exprs[0], + )? + .ok_or_else(err)?; + + // Collect the sorted expressions + let mut sorted_exprs = + vec![left_temp_sorted_filter_expr, right_temp_sorted_filter_expr]; + + // Build the expression interval graph + let mut graph = + ExprIntervalGraph::try_new(Arc::clone(filter.expression()), filter.schema())?; + + // Update sorted expressions with node indices + update_sorted_exprs_with_node_indices(&mut graph, &mut sorted_exprs); + + // Swap and remove to get the final sorted filter expressions + let right_sorted_filter_expr = sorted_exprs.swap_remove(1); + let left_sorted_filter_expr = sorted_exprs.swap_remove(0); + + Ok((left_sorted_filter_expr, right_sorted_filter_expr, graph)) +} + +#[cfg(test)] +pub mod tests { + + use super::*; + use crate::{joins::test_utils::complicated_filter, joins::utils::ColumnIndex}; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field}; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{binary, cast, col}; + + #[test] + fn test_column_exchange() -> Result<()> { + let left_child_schema = + Schema::new(vec![Field::new("left_1", DataType::Int32, true)]); + // Sorting information for the left side: + let left_child_sort_expr = PhysicalSortExpr { + expr: col("left_1", &left_child_schema)?, + options: SortOptions::default(), + }; + + let right_child_schema = Schema::new(vec![ + Field::new("right_1", DataType::Int32, true), + Field::new("right_2", DataType::Int32, true), + ]); + // Sorting information for the right side: + let right_child_sort_expr = PhysicalSortExpr { + expr: binary( + col("right_1", &right_child_schema)?, + Operator::Plus, + col("right_2", &right_child_schema)?, + &right_child_schema, + )?, + options: SortOptions::default(), + }; + + let intermediate_schema = Schema::new(vec![ + Field::new("filter_1", DataType::Int32, true), + Field::new("filter_2", DataType::Int32, true), + Field::new("filter_3", DataType::Int32, true), + ]); + // Our filter expression is: left_1 > right_1 + right_2. + let filter_left = col("filter_1", &intermediate_schema)?; + let filter_right = binary( + col("filter_2", &intermediate_schema)?, + Operator::Plus, + col("filter_3", &intermediate_schema)?, + &intermediate_schema, + )?; + let filter_expr = binary( + Arc::clone(&filter_left), + Operator::Gt, + Arc::clone(&filter_right), + &intermediate_schema, + )?; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ColumnIndex { + index: 1, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + let left_sort_filter_expr = build_filter_input_order( + JoinSide::Left, + &filter, + &Arc::new(left_child_schema), + &left_child_sort_expr, + )? + .unwrap(); + assert!(left_child_sort_expr.eq(left_sort_filter_expr.origin_sorted_expr())); + + let right_sort_filter_expr = build_filter_input_order( + JoinSide::Right, + &filter, + &Arc::new(right_child_schema), + &right_child_sort_expr, + )? + .unwrap(); + assert!(right_child_sort_expr.eq(right_sort_filter_expr.origin_sorted_expr())); + + // Assert that adjusted (left) filter expression matches with `left_child_sort_expr`: + assert!(filter_left.eq(left_sort_filter_expr.filter_expr())); + // Assert that adjusted (right) filter expression matches with `right_child_sort_expr`: + assert!(filter_right.eq(right_sort_filter_expr.filter_expr())); + Ok(()) + } + + #[test] + fn test_column_collector() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + Field::new("2", DataType::Int32, true), + ]); + let filter_expr = complicated_filter(&schema)?; + let columns = collect_columns(&filter_expr); + assert_eq!(columns.len(), 3); + Ok(()) + } + + #[test] + fn find_expr_inside_expr() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + Field::new("2", DataType::Int32, true), + ]); + let filter_expr = complicated_filter(&schema)?; + + let expr_1 = Arc::new(Column::new("gnz", 0)) as _; + assert!(!check_filter_expr_contains_sort_information( + &filter_expr, + &expr_1 + )); + + let expr_2 = col("1", &schema)? as _; + + assert!(check_filter_expr_contains_sort_information( + &filter_expr, + &expr_2 + )); + + let expr_3 = cast( + binary( + col("0", &schema)?, + Operator::Plus, + col("1", &schema)?, + &schema, + )?, + &schema, + DataType::Int64, + )?; + + assert!(check_filter_expr_contains_sort_information( + &filter_expr, + &expr_3 + )); + + let expr_4 = Arc::new(Column::new("1", 42)) as _; + + assert!(!check_filter_expr_contains_sort_information( + &filter_expr, + &expr_4, + )); + Ok(()) + } + + #[test] + fn build_sorted_expr() -> Result<()> { + let left_schema = Schema::new(vec![ + Field::new("la1", DataType::Int32, false), + Field::new("lb1", DataType::Int32, false), + Field::new("lc1", DataType::Int32, false), + Field::new("lt1", DataType::Int32, false), + Field::new("la2", DataType::Int32, false), + Field::new("la1_des", DataType::Int32, false), + ]); + + let right_schema = Schema::new(vec![ + Field::new("ra1", DataType::Int32, false), + Field::new("rb1", DataType::Int32, false), + Field::new("rc1", DataType::Int32, false), + Field::new("rt1", DataType::Int32, false), + Field::new("ra2", DataType::Int32, false), + Field::new("ra1_des", DataType::Int32, false), + ]); + + let intermediate_schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + Field::new("2", DataType::Int32, true), + ]); + let filter_expr = complicated_filter(&intermediate_schema)?; + let column_indices = vec![ + ColumnIndex { + index: left_schema.index_of("la1")?, + side: JoinSide::Left, + }, + ColumnIndex { + index: left_schema.index_of("la2")?, + side: JoinSide::Left, + }, + ColumnIndex { + index: right_schema.index_of("ra1")?, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + let left_schema = Arc::new(left_schema); + let right_schema = Arc::new(right_schema); + + assert!( + build_filter_input_order( + JoinSide::Left, + &filter, + &left_schema, + &PhysicalSortExpr { + expr: col("la1", left_schema.as_ref())?, + options: SortOptions::default(), + } + )? + .is_some() + ); + assert!( + build_filter_input_order( + JoinSide::Left, + &filter, + &left_schema, + &PhysicalSortExpr { + expr: col("lt1", left_schema.as_ref())?, + options: SortOptions::default(), + } + )? + .is_none() + ); + assert!( + build_filter_input_order( + JoinSide::Right, + &filter, + &right_schema, + &PhysicalSortExpr { + expr: col("ra1", right_schema.as_ref())?, + options: SortOptions::default(), + } + )? + .is_some() + ); + assert!( + build_filter_input_order( + JoinSide::Right, + &filter, + &right_schema, + &PhysicalSortExpr { + expr: col("rb1", right_schema.as_ref())?, + options: SortOptions::default(), + } + )? + .is_none() + ); + + Ok(()) + } + + // Test the case when we have an "ORDER BY a + b", and join filter condition includes "a - b". + #[test] + fn sorted_filter_expr_build() -> Result<()> { + let intermediate_schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + ]); + let filter_expr = binary( + col("0", &intermediate_schema)?, + Operator::Minus, + col("1", &intermediate_schema)?, + &intermediate_schema, + )?; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + + let sorted = PhysicalSortExpr { + expr: binary( + col("a", &schema)?, + Operator::Plus, + col("b", &schema)?, + &schema, + )?, + options: SortOptions::default(), + }; + + let res = convert_sort_expr_with_filter_schema( + &JoinSide::Left, + &filter, + &Arc::new(schema), + &sorted, + )?; + assert!(res.is_none()); + Ok(()) + } + + #[test] + fn test_shrink_if_necessary() { + let scale_factor = 4; + let mut join_hash_map = PruningJoinHashMap::with_capacity(100); + let data_size = 2000; + let deleted_part = 3 * data_size / 4; + // Add elements to the JoinHashMap + for hash_value in 0..data_size { + join_hash_map.map.insert_unique( + hash_value, + (hash_value, hash_value), + |(hash, _)| *hash, + ); + } + + assert_eq!(join_hash_map.map.len(), data_size as usize); + assert!(join_hash_map.map.capacity() >= data_size as usize); + + // Remove some elements from the JoinHashMap + for hash_value in 0..deleted_part { + join_hash_map + .map + .find_entry(hash_value, |(hash, _)| hash_value == *hash) + .unwrap() + .remove(); + } + + assert_eq!(join_hash_map.map.len(), (data_size - deleted_part) as usize); + + // Old capacity + let old_capacity = join_hash_map.map.capacity(); + + // Test shrink_if_necessary + join_hash_map.shrink_if_necessary(scale_factor); + + // The capacity should be reduced by the scale factor + let new_expected_capacity = + join_hash_map.map.capacity() * (scale_factor - 1) / scale_factor; + assert!(join_hash_map.map.capacity() >= new_expected_capacity); + assert!(join_hash_map.map.capacity() <= old_capacity); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/symmetric_hash_join.rs b/native/vendor/datafusion-physical-plan/src/joins/symmetric_hash_join.rs new file mode 100644 index 0000000..0c6e84b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/symmetric_hash_join.rs
@@ -0,0 +1,3035 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This file implements the symmetric hash join algorithm with range-based +//! data pruning to join two (potentially infinite) streams. +//! +//! A [`SymmetricHashJoinExec`] plan takes two children plan (with appropriate +//! output ordering) and produces the join output according to the given join +//! type and other options. +//! +//! This plan uses the [`OneSideHashJoiner`] object to facilitate join calculations +//! for both its children. + +use std::fmt::{self, Debug}; +use std::mem::{size_of, size_of_val}; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::vec; + +use crate::common::SharedMemoryReservation; +use crate::execution_plan::{boundedness_from_children, emission_type_from_children}; +use crate::joins::stream_join_utils::{ + PruningJoinHashMap, SortedFilterExpr, StreamJoinMetrics, + calculate_filter_expr_intervals, combine_two_batches, + convert_sort_expr_with_filter_schema, get_pruning_anti_indices, + get_pruning_semi_indices, prepare_sorted_exprs, record_visited_indices, +}; +use crate::joins::utils::{ + BatchSplitter, BatchTransformer, ColumnIndex, JoinFilter, JoinHashMapType, JoinOn, + JoinOnRef, NoopBatchTransformer, StatefulStreamResult, apply_join_filter_to_indices, + build_batch_from_indices, build_join_schema, check_join_is_valid, equal_rows_arr, + matchable_join_keys, symmetric_join_output_partitioning, update_hash, +}; +use crate::projection::{ + JoinData, ProjectionExec, try_pushdown_through_join_with_column_indices, +}; +use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, + joins::StreamJoinPartitionMode, + metrics::{ExecutionPlanMetricsSet, MetricsSet}, +}; + +use arrow::array::{ + ArrowPrimitiveType, NativeAdapter, PrimitiveArray, PrimitiveBuilder, UInt32Array, + UInt64Array, +}; +use arrow::compute::concat_batches; +use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::hash_utils::create_hashes; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::bisect; +use datafusion_common::{ + HashSet, JoinSide, JoinType, NullEquality, Result, assert_eq_or_internal_err, + plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_physical_expr::equivalence::join_equivalence_properties; +use datafusion_physical_expr::intervals::cp_solver::ExprIntervalGraph; +use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; + +use datafusion_common::hash_utils::RandomState; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::{Stream, StreamExt, ready}; + +const HASHMAP_SHRINK_SCALE_FACTOR: usize = 4; + +/// A symmetric hash join with range conditions is when both streams are hashed on the +/// join key and the resulting hash tables are used to join the streams. +/// The join is considered symmetric because the hash table is built on the join keys from both +/// streams, and the matching of rows is based on the values of the join keys in both streams. +/// This type of join is efficient in streaming context as it allows for fast lookups in the hash +/// table, rather than having to scan through one or both of the streams to find matching rows, also it +/// only considers the elements from the stream that fall within a certain sliding window (w/ range conditions), +/// making it more efficient and less likely to store stale data. This enables operating on unbounded streaming +/// data without any memory issues. +/// +/// For each input stream, create a hash table. +/// - For each new [RecordBatch] in build side, hash and insert into inputs hash table. Update offsets. +/// - Test if input is equal to a predefined set of other inputs. +/// - If so record the visited rows. If the matched row results must be produced (INNER, LEFT), output the [RecordBatch]. +/// - Try to prune other side (probe) with new [RecordBatch]. +/// - If the join type indicates that the unmatched rows results must be produced (LEFT, FULL etc.), +/// output the [RecordBatch] when a pruning happens or at the end of the data. +/// +/// +/// ``` text +/// +-------------------------+ +/// | | +/// left stream ---------| Left OneSideHashJoiner |---+ +/// | | | +/// +-------------------------+ | +/// | +/// |--------- Joined output +/// | +/// +-------------------------+ | +/// | | | +/// right stream ---------| Right OneSideHashJoiner |---+ +/// | | +/// +-------------------------+ +/// +/// Prune build side when the new RecordBatch comes to the probe side. We utilize interval arithmetic +/// on JoinFilter's sorted PhysicalExprs to calculate the joinable range. +/// +/// +/// PROBE SIDE BUILD SIDE +/// BUFFER BUFFER +/// +-------------+ +------------+ +/// | | | | Unjoinable +/// | | | | Range +/// | | | | +/// | | |--------------------------------- +/// | | | | | +/// | | | | | +/// | | / | | +/// | | | | | +/// | | | | | +/// | | | | | +/// | | | | | +/// | | | | | Joinable +/// | |/ | | Range +/// | || | | +/// |+-----------+|| | | +/// || Record || | | +/// || Batch || | | +/// |+-----------+|| | | +/// +-------------+\ +------------+ +/// | +/// \ +/// |--------------------------------- +/// +/// This happens when range conditions are provided on sorted columns. E.g. +/// +/// SELECT * FROM left_table, right_table +/// ON +/// left_key = right_key AND +/// left_time > right_time - INTERVAL 12 MINUTES AND left_time < right_time + INTERVAL 2 HOUR +/// +/// or +/// SELECT * FROM left_table, right_table +/// ON +/// left_key = right_key AND +/// left_sorted > right_sorted - 3 AND left_sorted < right_sorted + 10 +/// +/// For general purpose, in the second scenario, when the new data comes to probe side, the conditions can be used to +/// determine a specific threshold for discarding rows from the inner buffer. For example, if the sort order the +/// two columns ("left_sorted" and "right_sorted") are ascending (it can be different in another scenarios) +/// and the join condition is "left_sorted > right_sorted - 3" and the latest value on the right input is 1234, meaning +/// that the left side buffer must only keep rows where "leftTime > rightTime - 3 > 1234 - 3 > 1231" , +/// making the smallest value in 'left_sorted' 1231 and any rows below (since ascending) +/// than that can be dropped from the inner buffer. +/// ``` +#[derive(Debug, Clone)] +pub struct SymmetricHashJoinExec { + /// Left side stream + pub(crate) left: Arc<dyn ExecutionPlan>, + /// Right side stream + pub(crate) right: Arc<dyn ExecutionPlan>, + /// Set of common columns used to join on + pub(crate) on: Vec<(PhysicalExprRef, PhysicalExprRef)>, + /// Filters applied when finding matching rows + pub(crate) filter: Option<JoinFilter>, + /// How the join is performed + pub(crate) join_type: JoinType, + /// Shares the `RandomState` for the hashing algorithm + random_state: RandomState, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Information of index and left / right placement of columns + column_indices: Vec<ColumnIndex>, + /// Defines the null equality for the join. + pub(crate) null_equality: NullEquality, + /// Left side sort expression(s) + pub(crate) left_sort_exprs: Option<LexOrdering>, + /// Right side sort expression(s) + pub(crate) right_sort_exprs: Option<LexOrdering>, + /// Partition Mode + mode: StreamJoinPartitionMode, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl SymmetricHashJoinExec { + /// Tries to create a new [SymmetricHashJoinExec]. + /// # Error + /// This function errors when: + /// - It is not possible to join the left and right sides on keys `on`, or + /// - It fails to construct `SortedFilterExpr`s, or + /// - It fails to create the [ExprIntervalGraph]. + #[expect(clippy::too_many_arguments)] + pub fn try_new( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: Option<JoinFilter>, + join_type: &JoinType, + null_equality: NullEquality, + left_sort_exprs: Option<LexOrdering>, + right_sort_exprs: Option<LexOrdering>, + mode: StreamJoinPartitionMode, + ) -> Result<Self> { + let left_schema = left.schema(); + let right_schema = right.schema(); + + // Error out if no "on" constraints are given: + if on.is_empty() { + return plan_err!( + "On constraints in SymmetricHashJoinExec should be non-empty" + ); + } + + // Check if the join is valid with the given on constraints: + check_join_is_valid(&left_schema, &right_schema, &on)?; + + // Build the join schema from the left and right schemas: + let (schema, column_indices) = + build_join_schema(&left_schema, &right_schema, join_type); + + // Initialize the random state for the join operation: + let random_state = RandomState::with_seed(0); + let schema = Arc::new(schema); + let cache = Self::compute_properties(&left, &right, schema, *join_type, &on)?; + Ok(SymmetricHashJoinExec { + left, + right, + on, + filter, + join_type: *join_type, + random_state, + metrics: ExecutionPlanMetricsSet::new(), + column_indices, + null_equality, + left_sort_exprs, + right_sort_exprs, + mode, + cache: Arc::new(cache), + }) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + schema: SchemaRef, + join_type: JoinType, + join_on: JoinOnRef, + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let eq_properties = join_equivalence_properties( + left.equivalence_properties().clone(), + right.equivalence_properties().clone(), + &join_type, + schema, + &[false, false], + // Has alternating probe side + None, + join_on, + )?; + + let output_partitioning = + symmetric_join_output_partitioning(left, right, &join_type)?; + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + emission_type_from_children([left, right]), + boundedness_from_children([left, right]), + )) + } + + /// left stream + pub fn left(&self) -> &Arc<dyn ExecutionPlan> { + &self.left + } + + /// right stream + pub fn right(&self) -> &Arc<dyn ExecutionPlan> { + &self.right + } + + /// Set of common columns used to join on + pub fn on(&self) -> &[(PhysicalExprRef, PhysicalExprRef)] { + &self.on + } + + /// Filters applied before join output + pub fn filter(&self) -> Option<&JoinFilter> { + self.filter.as_ref() + } + + /// How the join is performed + pub fn join_type(&self) -> &JoinType { + &self.join_type + } + + /// Get null_equality + pub fn null_equality(&self) -> NullEquality { + self.null_equality + } + + /// Get partition mode + pub fn partition_mode(&self) -> StreamJoinPartitionMode { + self.mode + } + + /// Get left_sort_exprs + pub fn left_sort_exprs(&self) -> Option<&LexOrdering> { + self.left_sort_exprs.as_ref() + } + + /// Get right_sort_exprs + pub fn right_sort_exprs(&self) -> Option<&LexOrdering> { + self.right_sort_exprs.as_ref() + } + + /// Check if order information covers every column in the filter expression. + pub fn check_if_order_information_available(&self) -> Result<bool> { + if let Some(filter) = self.filter() { + let left = self.left(); + if let Some(left_ordering) = left.output_ordering() { + let right = self.right(); + if let Some(right_ordering) = right.output_ordering() { + let left_convertible = convert_sort_expr_with_filter_schema( + &JoinSide::Left, + filter, + &left.schema(), + &left_ordering[0], + )? + .is_some(); + let right_convertible = convert_sort_expr_with_filter_schema( + &JoinSide::Right, + filter, + &right.schema(), + &right_ordering[0], + )? + .is_some(); + return Ok(left_convertible && right_convertible); + } + } + } + Ok(false) + } +} + +impl DisplayAs for SymmetricHashJoinExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let display_filter = self.filter.as_ref().map_or_else( + || "".to_string(), + |f| format!(", filter={}", f.expression()), + ); + let on = self + .on + .iter() + .map(|(c1, c2)| format!("({c1}, {c2})")) + .collect::<Vec<String>>() + .join(", "); + write!( + f, + "SymmetricHashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}", + self.mode, self.join_type, on, display_filter + ) + } + DisplayFormatType::TreeRender => { + let on = self + .on + .iter() + .map(|(c1, c2)| { + format!("({} = {})", fmt_sql(c1.as_ref()), fmt_sql(c2.as_ref())) + }) + .collect::<Vec<String>>() + .join(", "); + + writeln!(f, "mode={:?}", self.mode)?; + if *self.join_type() != JoinType::Inner { + writeln!(f, "join_type={:?}", self.join_type)?; + } + writeln!(f, "on={on}") + } + } + } +} + +impl ExecutionPlan for SymmetricHashJoinExec { + fn name(&self) -> &'static str { + "SymmetricHashJoinExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + match self.mode { + StreamJoinPartitionMode::Partitioned => { + let (left_expr, right_expr) = self + .on + .iter() + .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) + .unzip(); + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) + } + StreamJoinPartitionMode::SinglePartition => { + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) + } + } + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![ + self.left_sort_exprs + .as_ref() + .map(|e| OrderingRequirements::from(e.clone())), + self.right_sort_exprs + .as_ref() + .map(|e| OrderingRequirements::from(e.clone())), + ] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.left, &self.right] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn crate::PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + let filter = self.filter.iter().map(|filter| filter.expression()); + crate::apply_expression_roots(join_keys.chain(filter), f) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(SymmetricHashJoinExec::try_new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.on.clone(), + self.filter.clone(), + &self.join_type, + self.null_equality, + self.left_sort_exprs.clone(), + self.right_sort_exprs.clone(), + self.mode, + )?)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let left_partitions = self.left.output_partitioning().partition_count(); + let right_partitions = self.right.output_partitioning().partition_count(); + assert_eq_or_internal_err!( + left_partitions, + right_partitions, + "Invalid SymmetricHashJoinExec, partition count mismatch {left_partitions}!={right_partitions},\ + consider using RepartitionExec" + ); + // If `filter_state` and `filter` are both present, then calculate sorted + // filter expressions for both sides, and build an expression graph. + let (left_sorted_filter_expr, right_sorted_filter_expr, graph) = match ( + self.left_sort_exprs(), + self.right_sort_exprs(), + &self.filter, + ) { + (Some(left_sort_exprs), Some(right_sort_exprs), Some(filter)) => { + let (left, right, graph) = prepare_sorted_exprs( + filter, + &self.left, + &self.right, + left_sort_exprs, + right_sort_exprs, + )?; + (Some(left), Some(right), Some(graph)) + } + // If `filter_state` or `filter` is not present, then return None + // for all three values: + _ => (None, None, None), + }; + + let (on_left, on_right) = self.on.iter().cloned().unzip(); + + let left_side_joiner = + OneSideHashJoiner::new(JoinSide::Left, on_left, self.left.schema()); + let right_side_joiner = + OneSideHashJoiner::new(JoinSide::Right, on_right, self.right.schema()); + + let left_stream = self.left.execute(partition, Arc::clone(&context))?; + + let right_stream = self.right.execute(partition, Arc::clone(&context))?; + + let batch_size = context.session_config().batch_size(); + let enforce_batch_size_in_joins = + context.session_config().enforce_batch_size_in_joins(); + + let reservation = Arc::new( + MemoryConsumer::new(format!("SymmetricHashJoinStream[{partition}]")) + .register(context.memory_pool()), + ); + if let Some(g) = graph.as_ref() { + reservation.try_grow(g.size())?; + } + + if enforce_batch_size_in_joins { + Ok(Box::pin(SymmetricHashJoinStream { + left_stream, + right_stream, + schema: self.schema(), + filter: self.filter.clone(), + join_type: self.join_type, + random_state: self.random_state.clone(), + left: left_side_joiner, + right: right_side_joiner, + column_indices: self.column_indices.clone(), + metrics: StreamJoinMetrics::new(partition, &self.metrics), + graph, + left_sorted_filter_expr, + right_sorted_filter_expr, + null_equality: self.null_equality, + state: SHJStreamState::PullRight, + reservation, + batch_transformer: BatchSplitter::new(batch_size), + })) + } else { + Ok(Box::pin(SymmetricHashJoinStream { + left_stream, + right_stream, + schema: self.schema(), + filter: self.filter.clone(), + join_type: self.join_type, + random_state: self.random_state.clone(), + left: left_side_joiner, + right: right_side_joiner, + column_indices: self.column_indices.clone(), + metrics: StreamJoinMetrics::new(partition, &self.metrics), + graph, + left_sorted_filter_expr, + right_sorted_filter_expr, + null_equality: self.null_equality, + state: SHJStreamState::PullRight, + reservation, + batch_transformer: NoopBatchTransformer::new(), + })) + } + } + + /// Tries to swap the projection with its input [`SymmetricHashJoinExec`]. If it can be done, + /// it returns the new swapped version having the [`SymmetricHashJoinExec`] as the top plan. + /// Otherwise, it returns None. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + let schema = self.schema(); + if let Some(JoinData { + projected_left_child, + projected_right_child, + join_filter, + join_on, + }) = try_pushdown_through_join_with_column_indices( + projection, + self.left(), + self.right(), + self.on(), + &schema, + self.filter(), + self.column_indices.as_slice(), + )? { + SymmetricHashJoinExec::try_new( + Arc::new(projected_left_child), + Arc::new(projected_right_child), + join_on, + join_filter, + self.join_type(), + self.null_equality(), + self.right().output_ordering().cloned(), + self.left().output_ordering().cloned(), + self.partition_mode(), + ) + .map(|e| Some(Arc::new(e) as _)) + } else { + Ok(None) + } + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + let on = self + .on() + .iter() + .map(|(left, right)| { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(left)?), + right: Some(ctx.encode_expr(right)?), + }) + }) + .collect::<Result<Vec<_>>>()?; + + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + }; + let null_equality = match self.null_equality() { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + }; + let partition_mode = match self.partition_mode() { + StreamJoinPartitionMode::SinglePartition => { + protobuf::StreamPartitionMode::SinglePartition + } + StreamJoinPartitionMode::Partitioned => { + protobuf::StreamPartitionMode::PartitionedExec + } + }; + let filter = self + .filter() + .map(|filter| -> Result<protobuf::JoinFilter> { + let expression = ctx.encode_expr(filter.expression())?; + let column_indices = filter + .column_indices() + .iter() + .map(|column_index| { + let side = match column_index.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: column_index.index as u32, + side: side.into(), + } + }) + .collect(); + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(filter.schema().as_ref().try_into()?), + }) + }) + .transpose()?; + let expr_ctx = ctx.expr_ctx(); + let left_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto( + self.left_sort_exprs(), + &expr_ctx, + )?; + let right_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto( + self.right_sort_exprs(), + &expr_ctx, + )?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin( + Box::new(protobuf::SymmetricHashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + left_sort_exprs, + right_sort_exprs, + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SymmetricHashJoinExec { + /// Reconstruct a [`SymmetricHashJoinExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_common::internal_datafusion_err; + use datafusion_proto_models::protobuf; + + let sym_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin, + "SymmetricHashJoinExec", + ); + let left = ctx.decode_required_child( + sym_join.left.as_deref(), + "SymmetricHashJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sym_join.right.as_deref(), + "SymmetricHashJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = sym_join + .on + .iter() + .map(|columns| { + let left = ctx.decode_required_expr( + columns.left.as_ref(), + left_schema.as_ref(), + "SymmetricHashJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + columns.right.as_ref(), + right_schema.as_ref(), + "SymmetricHashJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::<Result<JoinOn>>()?; + + let join_type = + match protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinType {}", + sym_join.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + let null_equality = match protobuf::NullEquality::try_from(sym_join.null_equality) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown NullEquality {}", + sym_join.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + let partition_mode = + match protobuf::StreamPartitionMode::try_from(sym_join.partition_mode) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown StreamPartitionMode {}", + sym_join.partition_mode + ) + })? { + protobuf::StreamPartitionMode::SinglePartition => { + StreamJoinPartitionMode::SinglePartition + } + protobuf::StreamPartitionMode::PartitionedExec => { + StreamJoinPartitionMode::Partitioned + } + }; + let filter = sym_join + .filter + .as_ref() + .map(|filter| -> Result<JoinFilter> { + let schema: Schema = filter + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "SymmetricHashJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + filter.expression.as_ref(), + &schema, + "SymmetricHashJoinExec", + "filter.expression", + )?; + let column_indices = filter + .column_indices + .iter() + .map(|column_index| { + let side = protobuf::JoinSide::try_from(column_index.side) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinSide {}", + column_index.side + ) + })?; + let side = match side { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: column_index.index as usize, + side, + }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + let left_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto( + &sym_join.left_sort_exprs, + &ctx.expr_ctx(left_schema.as_ref()), + )?; + let right_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto( + &sym_join.right_sort_exprs, + &ctx.expr_ctx(right_schema.as_ref()), + )?; + + Self::try_new( + left, + right, + on, + filter, + &join_type, + null_equality, + left_sort_exprs, + right_sort_exprs, + partition_mode, + ) + .map(|exec| Arc::new(exec) as _) + } +} + +/// A stream that issues [RecordBatch]es as they arrive from the right of the join. +struct SymmetricHashJoinStream<T> { + /// Input streams + left_stream: SendableRecordBatchStream, + right_stream: SendableRecordBatchStream, + /// Input schema + schema: Arc<Schema>, + /// join filter + filter: Option<JoinFilter>, + /// type of the join + join_type: JoinType, + // left hash joiner + left: OneSideHashJoiner, + /// right hash joiner + right: OneSideHashJoiner, + /// Information of index and left / right placement of columns + column_indices: Vec<ColumnIndex>, + // Expression graph for range pruning. + graph: Option<ExprIntervalGraph>, + // Left globally sorted filter expr + left_sorted_filter_expr: Option<SortedFilterExpr>, + // Right globally sorted filter expr + right_sorted_filter_expr: Option<SortedFilterExpr>, + /// Random state used for hashing initialization + random_state: RandomState, + /// Defines the null equality for the join. + null_equality: NullEquality, + /// Metrics + metrics: StreamJoinMetrics, + /// Memory reservation + reservation: SharedMemoryReservation, + /// State machine for input execution + state: SHJStreamState, + /// Transforms the output batch before returning. + batch_transformer: T, +} + +impl<T: BatchTransformer + Unpin + Send> RecordBatchStream + for SymmetricHashJoinStream<T> +{ + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl<T: BatchTransformer + Unpin + Send> Stream for SymmetricHashJoinStream<T> { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_impl(cx) + } +} + +/// Determine the pruning length for `buffer`. +/// +/// This function evaluates the build side filter expression, converts the +/// result into an array and determines the pruning length by performing a +/// binary search on the array. +/// +/// # Arguments +/// +/// * `buffer`: The record batch to be pruned. +/// * `build_side_filter_expr`: The filter expression on the build side used +/// to determine the pruning length. +/// +/// # Returns +/// +/// A [Result] object that contains the pruning length. The function will return +/// an error if +/// - there is an issue evaluating the build side filter expression; +/// - there is an issue converting the build side filter expression into an array +fn determine_prune_length( + buffer: &RecordBatch, + build_side_filter_expr: &SortedFilterExpr, +) -> Result<usize> { + let origin_sorted_expr = build_side_filter_expr.origin_sorted_expr(); + let interval = build_side_filter_expr.interval(); + // Evaluate the build side filter expression and convert it into an array + let batch_arr = origin_sorted_expr + .expr + .evaluate(buffer)? + .into_array(buffer.num_rows())?; + + // Get the lower or upper interval based on the sort direction + let target = if origin_sorted_expr.options.descending { + interval.upper().clone() + } else { + interval.lower().clone() + }; + + // Perform binary search on the array to determine the length of the record batch to be pruned + bisect::<true>(&[batch_arr], &[target], &[origin_sorted_expr.options]) +} + +/// This method determines if the result of the join should be produced in the final step or not. +/// +/// # Arguments +/// +/// * `build_side` - Enum indicating the side of the join used as the build side. +/// * `join_type` - Enum indicating the type of join to be performed. +/// +/// # Returns +/// +/// A boolean indicating whether the result of the join should be produced in the final step or not. +/// The result will be true if the build side is JoinSide::Left and the join type is one of +/// JoinType::Left, JoinType::LeftAnti, JoinType::Full or JoinType::LeftSemi. +/// If the build side is JoinSide::Right, the result will be true if the join type +/// is one of JoinType::Right, JoinType::RightAnti, JoinType::Full, or JoinType::RightSemi. +fn need_to_produce_result_in_final(build_side: JoinSide, join_type: JoinType) -> bool { + if build_side == JoinSide::Left { + matches!( + join_type, + JoinType::Left + | JoinType::LeftAnti + | JoinType::Full + | JoinType::LeftSemi + | JoinType::LeftMark + ) + } else { + matches!( + join_type, + JoinType::Right + | JoinType::RightAnti + | JoinType::Full + | JoinType::RightSemi + | JoinType::RightMark + ) + } +} + +/// Calculate indices by join type. +/// +/// This method returns a tuple of two arrays: build and probe indices. +/// The length of both arrays will be the same. +/// +/// # Arguments +/// +/// * `build_side`: Join side which defines the build side. +/// * `prune_length`: Length of the prune data. +/// * `visited_rows`: Hash set of visited rows of the build side. +/// * `deleted_offset`: Deleted offset of the build side. +/// * `join_type`: The type of join to be performed. +/// +/// # Returns +/// +/// A tuple of two arrays of primitive types representing the build and probe indices. +fn calculate_indices_by_join_type<L: ArrowPrimitiveType, R: ArrowPrimitiveType>( + build_side: JoinSide, + prune_length: usize, + visited_rows: &HashSet<usize>, + deleted_offset: usize, + join_type: JoinType, +) -> Result<(PrimitiveArray<L>, PrimitiveArray<R>)> +where + NativeAdapter<L>: From<<L as ArrowPrimitiveType>::Native>, +{ + // Store the result in a tuple + let result = match (build_side, join_type) { + // For a mark join we “mark” each build‐side row with a dummy 0 in the probe‐side index + // if it ever matched. For example, if + // + // prune_length = 5 + // deleted_offset = 0 + // visited_rows = {1, 3} + // + // then we produce: + // + // build_indices = [0, 1, 2, 3, 4] + // probe_indices = [None, Some(0), None, Some(0), None] + // + // Example: for each build row i in [0..5): + // – We always output its own index i in `build_indices` + // – We output `Some(0)` in `probe_indices[i]` if row i was ever visited, else `None` + (JoinSide::Left, JoinType::LeftMark) => { + let build_indices = (0..prune_length) + .map(L::Native::from_usize) + .collect::<PrimitiveArray<L>>(); + let probe_indices = (0..prune_length) + .map(|idx| { + // For mark join we output a dummy index 0 to indicate the row had a match + visited_rows + .contains(&(idx + deleted_offset)) + .then_some(R::Native::from_usize(0).unwrap()) + }) + .collect(); + (build_indices, probe_indices) + } + (JoinSide::Right, JoinType::RightMark) => { + let build_indices = (0..prune_length) + .map(L::Native::from_usize) + .collect::<PrimitiveArray<L>>(); + let probe_indices = (0..prune_length) + .map(|idx| { + // For mark join we output a dummy index 0 to indicate the row had a match + visited_rows + .contains(&(idx + deleted_offset)) + .then_some(R::Native::from_usize(0).unwrap()) + }) + .collect(); + (build_indices, probe_indices) + } + // In the case of `Left` or `Right` join, or `Full` join, get the anti indices + (JoinSide::Left, JoinType::Left | JoinType::LeftAnti) + | (JoinSide::Right, JoinType::Right | JoinType::RightAnti) + | (_, JoinType::Full) => { + let build_unmatched_indices = + get_pruning_anti_indices(prune_length, deleted_offset, visited_rows); + let mut builder = + PrimitiveBuilder::<R>::with_capacity(build_unmatched_indices.len()); + builder.append_nulls(build_unmatched_indices.len()); + let probe_indices = builder.finish(); + (build_unmatched_indices, probe_indices) + } + // In the case of `LeftSemi` or `RightSemi` join, get the semi indices + (JoinSide::Left, JoinType::LeftSemi) | (JoinSide::Right, JoinType::RightSemi) => { + let build_unmatched_indices = + get_pruning_semi_indices(prune_length, deleted_offset, visited_rows); + let mut builder = + PrimitiveBuilder::<R>::with_capacity(build_unmatched_indices.len()); + builder.append_nulls(build_unmatched_indices.len()); + let probe_indices = builder.finish(); + (build_unmatched_indices, probe_indices) + } + // The case of other join types is not considered + _ => unreachable!(), + }; + Ok(result) +} + +/// This function produces unmatched record results based on the build side, +/// join type and other parameters. +/// +/// The method uses first `prune_length` rows from the build side input buffer +/// to produce results. +/// +/// # Arguments +/// +/// * `output_schema` - The schema of the final output record batch. +/// * `prune_length` - The length of the determined prune length. +/// * `probe_schema` - The schema of the probe [RecordBatch]. +/// * `join_type` - The type of join to be performed. +/// * `column_indices` - Indices of columns that are being joined. +/// +/// # Returns +/// +/// * `Option<RecordBatch>` - The final output record batch if required, otherwise [None]. +pub(crate) fn build_side_determined_results( + build_hash_joiner: &OneSideHashJoiner, + output_schema: &SchemaRef, + prune_length: usize, + probe_schema: SchemaRef, + join_type: JoinType, + column_indices: &[ColumnIndex], +) -> Result<Option<RecordBatch>> { + // Check if we need to produce a result in the final output: + if prune_length > 0 + && need_to_produce_result_in_final(build_hash_joiner.build_side, join_type) + { + // Calculate the indices for build and probe sides based on join type and build side: + let (build_indices, probe_indices) = calculate_indices_by_join_type( + build_hash_joiner.build_side, + prune_length, + &build_hash_joiner.visited_rows, + build_hash_joiner.deleted_offset, + join_type, + )?; + + // Create an empty probe record batch: + let empty_probe_batch = RecordBatch::new_empty(probe_schema); + // Build the final result from the indices of build and probe sides: + build_batch_from_indices( + output_schema.as_ref(), + &build_hash_joiner.input_buffer, + &empty_probe_batch, + &build_indices, + &probe_indices, + column_indices, + build_hash_joiner.build_side, + join_type, + ) + .map(|batch| (batch.num_rows() > 0).then_some(batch)) + } else { + // If we don't need to produce a result, return None + Ok(None) + } +} + +/// This method performs a join between the build side input buffer and the probe side batch. +/// +/// # Arguments +/// +/// * `build_hash_joiner` - Build side hash joiner +/// * `probe_hash_joiner` - Probe side hash joiner +/// * `schema` - A reference to the schema of the output record batch. +/// * `join_type` - The type of join to be performed. +/// * `on_probe` - An array of columns on which the join will be performed. The columns are from the probe side of the join. +/// * `filter` - An optional filter on the join condition. +/// * `probe_batch` - The second record batch to be joined. +/// * `column_indices` - An array of columns to be selected for the result of the join. +/// * `random_state` - The random state for the join. +/// * `null_equality` - Indicates whether NULL values should be treated as equal when joining. +/// +/// # Returns +/// +/// A [Result] containing an optional record batch if the join type is not one of `LeftAnti`, `RightAnti`, `LeftSemi` or `RightSemi`. +/// If the join type is one of the above four, the function will return [None]. +#[expect(clippy::too_many_arguments)] +pub(crate) fn join_with_probe_batch( + build_hash_joiner: &mut OneSideHashJoiner, + probe_hash_joiner: &mut OneSideHashJoiner, + schema: &SchemaRef, + join_type: JoinType, + filter: Option<&JoinFilter>, + probe_batch: &RecordBatch, + column_indices: &[ColumnIndex], + random_state: &RandomState, + null_equality: NullEquality, +) -> Result<Option<RecordBatch>> { + if build_hash_joiner.input_buffer.num_rows() == 0 || probe_batch.num_rows() == 0 { + return Ok(None); + } + let (build_indices, probe_indices) = lookup_join_hashmap( + &build_hash_joiner.hashmap, + &build_hash_joiner.input_buffer, + probe_batch, + &build_hash_joiner.on, + &probe_hash_joiner.on, + random_state, + null_equality, + &mut build_hash_joiner.hashes_buffer, + Some(build_hash_joiner.deleted_offset), + )?; + + let (build_indices, probe_indices) = if let Some(filter) = filter { + apply_join_filter_to_indices( + &build_hash_joiner.input_buffer, + probe_batch, + build_indices, + probe_indices, + filter, + build_hash_joiner.build_side, + None, + join_type, + )? + } else { + (build_indices, probe_indices) + }; + + if need_to_produce_result_in_final(build_hash_joiner.build_side, join_type) { + record_visited_indices( + &mut build_hash_joiner.visited_rows, + build_hash_joiner.deleted_offset, + &build_indices, + ); + } + if need_to_produce_result_in_final(build_hash_joiner.build_side.negate(), join_type) { + record_visited_indices( + &mut probe_hash_joiner.visited_rows, + probe_hash_joiner.offset, + &probe_indices, + ); + } + if matches!( + join_type, + JoinType::LeftAnti + | JoinType::RightAnti + | JoinType::LeftSemi + | JoinType::LeftMark + | JoinType::RightSemi + | JoinType::RightMark + ) { + Ok(None) + } else { + build_batch_from_indices( + schema, + &build_hash_joiner.input_buffer, + probe_batch, + &build_indices, + &probe_indices, + column_indices, + build_hash_joiner.build_side, + join_type, + ) + .map(|batch| (batch.num_rows() > 0).then_some(batch)) + } +} + +/// This method performs lookups against JoinHashMap by hash values of join-key columns, and handles potential +/// hash collisions. +/// +/// # Arguments +/// +/// * `build_hashmap` - hashmap collected from build side data. +/// * `build_batch` - Build side record batch. +/// * `probe_batch` - Probe side record batch. +/// * `build_on` - An array of columns on which the join will be performed. The columns are from the build side of the join. +/// * `probe_on` - An array of columns on which the join will be performed. The columns are from the probe side of the join. +/// * `random_state` - The random state for the join. +/// * `null_equality` - Indicates whether NULL values should be treated as equal when joining. +/// * `hashes_buffer` - Buffer used for probe side keys hash calculation. +/// * `deleted_offset` - deleted offset for build side data. +/// +/// # Returns +/// +/// A [Result] containing a tuple with two equal length arrays, representing indices of rows from build and probe side, +/// matched by join key columns. +#[expect(clippy::too_many_arguments)] +fn lookup_join_hashmap( + build_hashmap: &PruningJoinHashMap, + build_batch: &RecordBatch, + probe_batch: &RecordBatch, + build_on: &[PhysicalExprRef], + probe_on: &[PhysicalExprRef], + random_state: &RandomState, + null_equality: NullEquality, + hashes_buffer: &mut Vec<u64>, + deleted_offset: Option<usize>, +) -> Result<(UInt64Array, UInt32Array)> { + let keys_values = evaluate_expressions_to_arrays(probe_on, probe_batch)?; + let build_join_values = evaluate_expressions_to_arrays(build_on, build_batch)?; + + hashes_buffer.clear(); + hashes_buffer.resize(probe_batch.num_rows(), 0); + let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?; + + // As SymmetricHashJoin uses LIFO JoinHashMap, the chained list algorithm + // will return build indices for each probe row in a reverse order as such: + // Build Indices: [5, 4, 3] + // Probe Indices: [1, 1, 1] + // + // This affects the output sequence. Hypothetically, it's possible to preserve the lexicographic order on the build side. + // Let's consider probe rows [0,1] as an example: + // + // When the probe iteration sequence is reversed, the following pairings can be derived: + // + // For probe row 1: + // (5, 1) + // (4, 1) + // (3, 1) + // + // For probe row 0: + // (5, 0) + // (4, 0) + // (3, 0) + // + // After reversing both sets of indices, we obtain reversed indices: + // + // (3,0) + // (4,0) + // (5,0) + // (3,1) + // (4,1) + // (5,1) + // + // With this approach, the lexicographic order on both the probe side and the build side is preserved. + // + // Probe rows whose key contains a NULL cannot match any build row and are + // skipped without a map lookup. + let valid_keys = matchable_join_keys(&keys_values, null_equality); + let (mut matched_probe, mut matched_build) = build_hashmap.get_matched_indices( + Box::new( + hash_values + .iter() + .enumerate() + .filter(|(i, _)| { + valid_keys.as_ref().is_none_or(|valid| valid.is_valid(*i)) + }) + .rev(), + ), + deleted_offset, + ); + + matched_probe.reverse(); + matched_build.reverse(); + + let build_indices: UInt64Array = matched_build.into(); + let probe_indices: UInt32Array = matched_probe.into(); + + let (build_indices, probe_indices) = equal_rows_arr( + &build_indices, + &probe_indices, + &build_join_values, + &keys_values, + null_equality, + )?; + + Ok((build_indices, probe_indices)) +} + +pub struct OneSideHashJoiner { + /// Build side + build_side: JoinSide, + /// Input record batch buffer + pub input_buffer: RecordBatch, + /// Columns from the side + pub(crate) on: Vec<PhysicalExprRef>, + /// Hashmap + pub(crate) hashmap: PruningJoinHashMap, + /// Reuse the hashes buffer + pub(crate) hashes_buffer: Vec<u64>, + /// Matched rows + pub(crate) visited_rows: HashSet<usize>, + /// Offset + pub(crate) offset: usize, + /// Deleted offset + pub(crate) deleted_offset: usize, +} + +impl OneSideHashJoiner { + pub fn size(&self) -> usize { + let mut size = 0; + size += size_of_val(self); + size += size_of_val(&self.build_side); + size += self.input_buffer.get_array_memory_size(); + size += size_of_val(&self.on); + size += self.hashmap.size(); + size += self.hashes_buffer.capacity() * size_of::<u64>(); + size += self.visited_rows.capacity() * size_of::<usize>(); + size += size_of_val(&self.offset); + size += size_of_val(&self.deleted_offset); + size + } + pub fn new( + build_side: JoinSide, + on: Vec<PhysicalExprRef>, + schema: SchemaRef, + ) -> Self { + Self { + build_side, + input_buffer: RecordBatch::new_empty(schema), + on, + hashmap: PruningJoinHashMap::with_capacity(0), + hashes_buffer: vec![], + visited_rows: HashSet::new(), + offset: 0, + deleted_offset: 0, + } + } + + /// Updates the internal state of the [OneSideHashJoiner] with the incoming batch. + /// + /// # Arguments + /// + /// * `batch` - The incoming [RecordBatch] to be merged with the internal input buffer + /// * `random_state` - The random state used to hash values + /// * `null_equality` - Null semantics to use + /// + /// # Returns + /// + /// Returns a [Result] encapsulating any intermediate errors. + pub(crate) fn update_internal_state( + &mut self, + batch: &RecordBatch, + random_state: &RandomState, + null_equality: NullEquality, + ) -> Result<()> { + // Merge the incoming batch with the existing input buffer: + self.input_buffer = concat_batches(&batch.schema(), [&self.input_buffer, batch])?; + // Resize the hashes buffer to the number of rows in the incoming batch: + self.hashes_buffer.resize(batch.num_rows(), 0); + // Get allocation_info before adding the item + // Update the hashmap with the join key values and hashes of the incoming batch: + update_hash( + &self.on, + batch, + &mut self.hashmap, + self.offset, + random_state, + &mut self.hashes_buffer, + self.deleted_offset, + false, + null_equality, + )?; + Ok(()) + } + + /// Calculate prune length. + /// + /// # Arguments + /// + /// * `build_side_sorted_filter_expr` - Build side mutable sorted filter expression.. + /// * `probe_side_sorted_filter_expr` - Probe side mutable sorted filter expression. + /// * `graph` - A mutable reference to the physical expression graph. + /// + /// # Returns + /// + /// A Result object that contains the pruning length. + pub(crate) fn calculate_prune_length_with_probe_batch( + &mut self, + build_side_sorted_filter_expr: &mut SortedFilterExpr, + probe_side_sorted_filter_expr: &mut SortedFilterExpr, + graph: &mut ExprIntervalGraph, + ) -> Result<usize> { + // Return early if the input buffer is empty: + if self.input_buffer.num_rows() == 0 { + return Ok(0); + } + // Process the build and probe side sorted filter expressions if both are present: + // Collect the sorted filter expressions into a vector of (node_index, interval) tuples: + let mut filter_intervals = vec![]; + for expr in [ + &build_side_sorted_filter_expr, + &probe_side_sorted_filter_expr, + ] { + filter_intervals.push((expr.node_index(), expr.interval().clone())) + } + // Update the physical expression graph using the join filter intervals: + graph.update_ranges(&mut filter_intervals, Interval::TRUE)?; + // Extract the new join filter interval for the build side: + let calculated_build_side_interval = filter_intervals.remove(0).1; + // If the intervals have not changed, return early without pruning: + if calculated_build_side_interval.eq(build_side_sorted_filter_expr.interval()) { + return Ok(0); + } + // Update the build side interval and determine the pruning length: + build_side_sorted_filter_expr.set_interval(calculated_build_side_interval); + + determine_prune_length(&self.input_buffer, build_side_sorted_filter_expr) + } + + pub(crate) fn prune_internal_state(&mut self, prune_length: usize) -> Result<()> { + // Prune the hash values: + self.hashmap.prune_hash_values( + prune_length, + self.deleted_offset as u64, + HASHMAP_SHRINK_SCALE_FACTOR, + ); + // Remove pruned rows from the visited rows set: + for row in self.deleted_offset..(self.deleted_offset + prune_length) { + self.visited_rows.remove(&row); + } + // Update the input buffer after pruning: + self.input_buffer = self + .input_buffer + .slice(prune_length, self.input_buffer.num_rows() - prune_length); + // Increment the deleted offset: + self.deleted_offset += prune_length; + Ok(()) + } +} + +/// `SymmetricHashJoinStream` manages incremental join operations between two +/// streams. Unlike traditional join approaches that need to scan one side of +/// the join fully before proceeding, `SymmetricHashJoinStream` facilitates +/// more dynamic join operations by working with streams as they emit data. This +/// approach allows for more efficient processing, particularly in scenarios +/// where waiting for complete data materialization is not feasible or optimal. +/// The trait provides a framework for handling various states of such a join +/// process, ensuring that join logic is efficiently executed as data becomes +/// available from either stream. +/// +/// This implementation performs eager joins of data from two different asynchronous +/// streams, typically referred to as left and right streams. The implementation +/// provides a comprehensive set of methods to control and execute the join +/// process, leveraging the states defined in `SHJStreamState`. Methods are +/// primarily focused on asynchronously fetching data batches from each stream, +/// processing them, and managing transitions between various states of the join. +/// +/// This implementations use a state machine approach to navigate different +/// stages of the join operation, handling data from both streams and determining +/// when the join completes. +/// +/// State Transitions: +/// - From `PullLeft` to `PullRight` or `LeftExhausted`: +/// - In `fetch_next_from_left_stream`, when fetching a batch from the left stream: +/// - On success (`Some(Ok(batch))`), state transitions to `PullRight` for +/// processing the batch. +/// - On error (`Some(Err(e))`), the error is returned, and the state remains +/// unchanged. +/// - On no data (`None`), state changes to `LeftExhausted`, returning `Continue` +/// to proceed with the join process. +/// - From `PullRight` to `PullLeft` or `RightExhausted`: +/// - In `fetch_next_from_right_stream`, when fetching from the right stream: +/// - If a batch is available, state changes to `PullLeft` for processing. +/// - On error, the error is returned without changing the state. +/// - If right stream is exhausted (`None`), state transitions to `RightExhausted`, +/// with a `Continue` result. +/// - Handling `RightExhausted` and `LeftExhausted`: +/// - Methods `handle_right_stream_end` and `handle_left_stream_end` manage scenarios +/// when streams are exhausted: +/// - They attempt to continue processing with the other stream. +/// - If both streams are exhausted, state changes to `BothExhausted { final_result: false }`. +/// - Transition to `BothExhausted { final_result: true }`: +/// - Occurs in `prepare_for_final_results_after_exhaustion` when both streams are +/// exhausted, indicating completion of processing and availability of final results. +impl<T: BatchTransformer> SymmetricHashJoinStream<T> { + /// Implements the main polling logic for the join stream. + /// + /// This method continuously checks the state of the join stream and + /// acts accordingly by delegating the handling to appropriate sub-methods + /// depending on the current state. + /// + /// # Arguments + /// + /// * `cx` - A context that facilitates cooperative non-blocking execution within a task. + /// + /// # Returns + /// + /// * `Poll<Option<Result<RecordBatch>>>` - A polled result, either a `RecordBatch` or None. + fn poll_next_impl( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + match self.batch_transformer.next() { + None => { + let result = match self.state() { + SHJStreamState::PullRight => { + ready!(self.fetch_next_from_right_stream(cx)) + } + SHJStreamState::PullLeft => { + ready!(self.fetch_next_from_left_stream(cx)) + } + SHJStreamState::RightExhausted => { + ready!(self.handle_right_stream_end(cx)) + } + SHJStreamState::LeftExhausted => { + ready!(self.handle_left_stream_end(cx)) + } + SHJStreamState::BothExhausted { + final_result: false, + } => self.prepare_for_final_results_after_exhaustion(), + SHJStreamState::BothExhausted { final_result: true } => { + return Poll::Ready(None); + } + }; + + match result? { + StatefulStreamResult::Ready(None) => { + return Poll::Ready(None); + } + StatefulStreamResult::Ready(Some(batch)) => { + self.batch_transformer.set_batch(batch); + } + _ => {} + } + } + Some((batch, _)) => { + return self + .metrics + .baseline_metrics + .record_poll(Poll::Ready(Some(Ok(batch)))); + } + } + } + } + + /// Release the right input pipeline's resources. + fn cleanup_depleted_right_stream(&mut self) { + let right_schema = self.right_stream.schema(); + self.right_stream = Box::pin(EmptyRecordBatchStream::new(right_schema)); + } + + /// Release the left input pipeline's resources. + fn cleanup_depleted_left_stream(&mut self) { + let left_schema = self.left_stream.schema(); + self.left_stream = Box::pin(EmptyRecordBatchStream::new(left_schema)); + } + + /// Asynchronously pulls the next batch from the right stream. + /// + /// This default implementation checks for the next value in the right stream. + /// If a batch is found, the state is switched to `PullLeft`, and the batch handling + /// is delegated to `process_batch_from_right`. If the stream ends, the state is set to `RightExhausted`. + /// + /// # Returns + /// + /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after pulling the batch. + fn fetch_next_from_right_stream( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.right_stream().poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() == 0 { + return Poll::Ready(Ok(StatefulStreamResult::Continue)); + } + self.set_state(SHJStreamState::PullLeft); + Poll::Ready(self.process_batch_from_right(&batch)) + } + Some(Err(e)) => Poll::Ready(Err(e)), + None => { + self.cleanup_depleted_right_stream(); + self.set_state(SHJStreamState::RightExhausted); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + } + } + + /// Asynchronously pulls the next batch from the left stream. + /// + /// This default implementation checks for the next value in the left stream. + /// If a batch is found, the state is switched to `PullRight`, and the batch handling + /// is delegated to `process_batch_from_left`. If the stream ends, the state is set to `LeftExhausted`. + /// + /// # Returns + /// + /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after pulling the batch. + fn fetch_next_from_left_stream( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.left_stream().poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() == 0 { + return Poll::Ready(Ok(StatefulStreamResult::Continue)); + } + self.set_state(SHJStreamState::PullRight); + Poll::Ready(self.process_batch_from_left(&batch)) + } + Some(Err(e)) => Poll::Ready(Err(e)), + None => { + self.cleanup_depleted_left_stream(); + self.set_state(SHJStreamState::LeftExhausted); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + } + } + + /// Asynchronously handles the scenario when the right stream is exhausted. + /// + /// In this default implementation, when the right stream is exhausted, it attempts + /// to pull from the left stream. If a batch is found in the left stream, it delegates + /// the handling to `process_batch_from_left`. If both streams are exhausted, the state is set + /// to indicate both streams are exhausted without final results yet. + /// + /// # Returns + /// + /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after checking the exhaustion state. + fn handle_right_stream_end( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.left_stream().poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() == 0 { + return Poll::Ready(Ok(StatefulStreamResult::Continue)); + } + Poll::Ready(self.process_batch_after_right_end(&batch)) + } + Some(Err(e)) => Poll::Ready(Err(e)), + None => { + self.cleanup_depleted_left_stream(); + self.set_state(SHJStreamState::BothExhausted { + final_result: false, + }); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + } + } + + /// Asynchronously handles the scenario when the left stream is exhausted. + /// + /// When the left stream is exhausted, this default + /// implementation tries to pull from the right stream and delegates the batch + /// handling to `process_batch_after_left_end`. If both streams are exhausted, the state + /// is updated to indicate so. + /// + /// # Returns + /// + /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after checking the exhaustion state. + fn handle_left_stream_end( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Result<StatefulStreamResult<Option<RecordBatch>>>> { + match ready!(self.right_stream().poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() == 0 { + return Poll::Ready(Ok(StatefulStreamResult::Continue)); + } + Poll::Ready(self.process_batch_after_left_end(&batch)) + } + Some(Err(e)) => Poll::Ready(Err(e)), + None => { + self.cleanup_depleted_right_stream(); + self.set_state(SHJStreamState::BothExhausted { + final_result: false, + }); + Poll::Ready(Ok(StatefulStreamResult::Continue)) + } + } + } + + /// Handles the state when both streams are exhausted and final results are yet to be produced. + /// + /// This default implementation switches the state to indicate both streams are + /// exhausted with final results and then invokes the handling for this specific + /// scenario via `process_batches_before_finalization`. + /// + /// # Returns + /// + /// * `Result<StatefulStreamResult<Option<RecordBatch>>>` - The state result after both streams are exhausted. + fn prepare_for_final_results_after_exhaustion( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + self.set_state(SHJStreamState::BothExhausted { final_result: true }); + self.process_batches_before_finalization() + } + + fn process_batch_from_right( + &mut self, + batch: &RecordBatch, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + self.perform_join_for_given_side(batch, JoinSide::Right) + .map(|maybe_batch| { + if maybe_batch.is_some() { + StatefulStreamResult::Ready(maybe_batch) + } else { + StatefulStreamResult::Continue + } + }) + } + + fn process_batch_from_left( + &mut self, + batch: &RecordBatch, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + self.perform_join_for_given_side(batch, JoinSide::Left) + .map(|maybe_batch| { + if maybe_batch.is_some() { + StatefulStreamResult::Ready(maybe_batch) + } else { + StatefulStreamResult::Continue + } + }) + } + + fn process_batch_after_left_end( + &mut self, + right_batch: &RecordBatch, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + self.process_batch_from_right(right_batch) + } + + fn process_batch_after_right_end( + &mut self, + left_batch: &RecordBatch, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + self.process_batch_from_left(left_batch) + } + + fn process_batches_before_finalization( + &mut self, + ) -> Result<StatefulStreamResult<Option<RecordBatch>>> { + // Get the left side results: + let left_result = build_side_determined_results( + &self.left, + &self.schema, + self.left.input_buffer.num_rows(), + self.right.input_buffer.schema(), + self.join_type, + &self.column_indices, + )?; + // Get the right side results: + let right_result = build_side_determined_results( + &self.right, + &self.schema, + self.right.input_buffer.num_rows(), + self.left.input_buffer.schema(), + self.join_type, + &self.column_indices, + )?; + + // Combine the left and right results: + let result = combine_two_batches(&self.schema, left_result, right_result)?; + + // Return the result: + if result.is_some() { + return Ok(StatefulStreamResult::Ready(result)); + } + Ok(StatefulStreamResult::Continue) + } + + fn right_stream(&mut self) -> &mut SendableRecordBatchStream { + &mut self.right_stream + } + + fn left_stream(&mut self) -> &mut SendableRecordBatchStream { + &mut self.left_stream + } + + fn set_state(&mut self, state: SHJStreamState) { + self.state = state; + } + + fn state(&mut self) -> SHJStreamState { + self.state.clone() + } + + fn size(&self) -> usize { + let mut size = 0; + size += size_of_val(&self.schema); + size += size_of_val(&self.filter); + size += size_of_val(&self.join_type); + size += self.left.size(); + size += self.right.size(); + size += size_of_val(&self.column_indices); + size += self.graph.as_ref().map(|g| g.size()).unwrap_or(0); + size += size_of_val(&self.left_sorted_filter_expr); + size += size_of_val(&self.right_sorted_filter_expr); + size += size_of_val(&self.random_state); + size += size_of_val(&self.null_equality); + size += size_of_val(&self.metrics); + size + } + + /// Performs a join operation for the specified `probe_side` (either left or right). + /// This function: + /// 1. Determines which side is the probe and which is the build side. + /// 2. Updates metrics based on the batch that was polled. + /// 3. Executes the join with the given `probe_batch`. + /// 4. Optionally computes anti-join results if all conditions are met. + /// 5. Combines the results and returns a combined batch or `None` if no batch was produced. + fn perform_join_for_given_side( + &mut self, + probe_batch: &RecordBatch, + probe_side: JoinSide, + ) -> Result<Option<RecordBatch>> { + let ( + probe_hash_joiner, + build_hash_joiner, + probe_side_sorted_filter_expr, + build_side_sorted_filter_expr, + probe_side_metrics, + ) = if probe_side.eq(&JoinSide::Left) { + ( + &mut self.left, + &mut self.right, + &mut self.left_sorted_filter_expr, + &mut self.right_sorted_filter_expr, + &mut self.metrics.left, + ) + } else { + ( + &mut self.right, + &mut self.left, + &mut self.right_sorted_filter_expr, + &mut self.left_sorted_filter_expr, + &mut self.metrics.right, + ) + }; + // Update the metrics for the stream that was polled: + probe_side_metrics.input_batches.add(1); + probe_side_metrics.input_rows.add(probe_batch.num_rows()); + // Update the internal state of the hash joiner for the build side: + probe_hash_joiner.update_internal_state( + probe_batch, + &self.random_state, + self.null_equality, + )?; + // Join the two sides: + let equal_result = join_with_probe_batch( + build_hash_joiner, + probe_hash_joiner, + &self.schema, + self.join_type, + self.filter.as_ref(), + probe_batch, + &self.column_indices, + &self.random_state, + self.null_equality, + )?; + // Increment the offset for the probe hash joiner: + probe_hash_joiner.offset += probe_batch.num_rows(); + + let anti_result = if let ( + Some(build_side_sorted_filter_expr), + Some(probe_side_sorted_filter_expr), + Some(graph), + ) = ( + build_side_sorted_filter_expr.as_mut(), + probe_side_sorted_filter_expr.as_mut(), + self.graph.as_mut(), + ) { + // Calculate filter intervals: + calculate_filter_expr_intervals( + &build_hash_joiner.input_buffer, + build_side_sorted_filter_expr, + probe_batch, + probe_side_sorted_filter_expr, + )?; + let prune_length = build_hash_joiner + .calculate_prune_length_with_probe_batch( + build_side_sorted_filter_expr, + probe_side_sorted_filter_expr, + graph, + )?; + let result = build_side_determined_results( + build_hash_joiner, + &self.schema, + prune_length, + probe_batch.schema(), + self.join_type, + &self.column_indices, + )?; + build_hash_joiner.prune_internal_state(prune_length)?; + result + } else { + None + }; + + // Combine results: + let result = combine_two_batches(&self.schema, equal_result, anti_result)?; + let capacity = self.size(); + self.metrics.stream_memory_usage.set(capacity); + self.reservation.try_resize(capacity)?; + Ok(result) + } +} + +/// Represents the various states of an symmetric hash join stream operation. +/// +/// This enum is used to track the current state of streaming during a join +/// operation. It provides indicators as to which side of the join needs to be +/// pulled next or if one (or both) sides have been exhausted. This allows +/// for efficient management of resources and optimal performance during the +/// join process. +#[derive(Clone, Debug)] +pub enum SHJStreamState { + /// Indicates that the next step should pull from the right side of the join. + PullRight, + + /// Indicates that the next step should pull from the left side of the join. + PullLeft, + + /// State representing that the right side of the join has been fully processed. + RightExhausted, + + /// State representing that the left side of the join has been fully processed. + LeftExhausted, + + /// Represents a state where both sides of the join are exhausted. + /// + /// The `final_result` field indicates whether the join operation has + /// produced a final result or not. + BothExhausted { final_result: bool }, +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::{LazyLock, Mutex}; + + use super::*; + use crate::joins::test_utils::{ + build_sides_record_batches, compare_batches, complicated_filter, + create_memory_table, join_expr_tests_fixture_f64, join_expr_tests_fixture_i32, + join_expr_tests_fixture_temporal, partitioned_hash_join_with_filter, + partitioned_sym_join_with_filter, split_record_batches, + }; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, IntervalUnit, TimeUnit}; + use datafusion_common::ScalarValue; + use datafusion_execution::config::SessionConfig; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{Column, binary, col, lit}; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + + use rstest::*; + + const TABLE_SIZE: i32 = 30; + + type TableKey = (i32, i32, usize); // (cardinality.0, cardinality.1, batch_size) + type TableValue = (Vec<RecordBatch>, Vec<RecordBatch>); // (left, right) + + // Cache for storing tables + static TABLE_CACHE: LazyLock<Mutex<HashMap<TableKey, TableValue>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + + fn get_or_create_table( + cardinality: (i32, i32), + batch_size: usize, + ) -> Result<TableValue> { + { + let cache = TABLE_CACHE.lock().unwrap(); + if let Some(table) = cache.get(&(cardinality.0, cardinality.1, batch_size)) { + return Ok(table.clone()); + } + } + + // If not, create the table + let (left_batch, right_batch) = + build_sides_record_batches(TABLE_SIZE, cardinality)?; + + let (left_partition, right_partition) = ( + split_record_batches(&left_batch, batch_size)?, + split_record_batches(&right_batch, batch_size)?, + ); + + // Lock the cache again and store the table + let mut cache = TABLE_CACHE.lock().unwrap(); + + // Store the table in the cache + cache.insert( + (cardinality.0, cardinality.1, batch_size), + (left_partition.clone(), right_partition.clone()), + ); + + Ok((left_partition, right_partition)) + } + + pub async fn experiment( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + filter: Option<JoinFilter>, + join_type: JoinType, + on: JoinOn, + task_ctx: Arc<TaskContext>, + ) -> Result<()> { + let first_batches = partitioned_sym_join_with_filter( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + filter.clone(), + &join_type, + NullEquality::NullEqualsNothing, + Arc::clone(&task_ctx), + ) + .await?; + let second_batches = partitioned_hash_join_with_filter( + left, + right, + on, + filter, + &join_type, + NullEquality::NullEqualsNothing, + task_ctx, + ) + .await?; + compare_batches(&first_batches, &second_batches); + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn complex_join_all_one_ascending_numeric( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values( + (4, 5), + (12, 17), + )] + cardinality: (i32, i32), + ) -> Result<()> { + // a + b > c + 10 AND a + b < c + 100 + let task_ctx = Arc::new(TaskContext::default()); + + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + + let left_sorted = [PhysicalSortExpr { + expr: binary( + col("la1", left_schema)?, + Operator::Plus, + col("la2", left_schema)?, + left_schema, + )?, + options: SortOptions::default(), + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("ra1", right_schema)?, + options: SortOptions::default(), + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![( + binary( + col("lc1", left_schema)?, + Operator::Plus, + lit(ScalarValue::Int32(Some(1))), + left_schema, + )?, + Arc::new(Column::new_with_schema("rc1", right_schema)?) as _, + )]; + + let intermediate_schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + Field::new("2", DataType::Int32, true), + ]); + let filter_expr = complicated_filter(&intermediate_schema)?; + let column_indices = vec![ + ColumnIndex { + index: left_schema.index_of("la1")?, + side: JoinSide::Left, + }, + ColumnIndex { + index: left_schema.index_of("la2")?, + side: JoinSide::Left, + }, + ColumnIndex { + index: right_schema.index_of("ra1")?, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn join_all_one_ascending_numeric( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values(0, 1, 2, 3, 4, 5)] case_expr: usize, + ) -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let (left_partition, right_partition) = get_or_create_table((4, 5), 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + + let left_sorted = [PhysicalSortExpr { + expr: col("la1", left_schema)?, + options: SortOptions::default(), + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("ra1", right_schema)?, + options: SortOptions::default(), + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + ]); + let filter_expr = join_expr_tests_fixture_i32( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn join_without_sort_information( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values(0, 1, 2, 3, 4, 5)] case_expr: usize, + ) -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let (left_partition, right_partition) = get_or_create_table((4, 5), 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let (left, right) = + create_memory_table(left_partition, right_partition, vec![], vec![])?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + ]); + let filter_expr = join_expr_tests_fixture_i32( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 5, + side: JoinSide::Left, + }, + ColumnIndex { + index: 5, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn join_without_filter( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + ) -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let (left_partition, right_partition) = get_or_create_table((11, 21), 8)?; + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let (left, right) = + create_memory_table(left_partition, right_partition, vec![], vec![])?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + experiment(left, right, None, join_type, on, task_ctx).await?; + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn join_all_one_descending_numeric_particular( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values(0, 1, 2, 3, 4, 5)] case_expr: usize, + ) -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let (left_partition, right_partition) = get_or_create_table((11, 21), 8)?; + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = [PhysicalSortExpr { + expr: col("la1_des", left_schema)?, + options: SortOptions { + descending: true, + nulls_first: true, + }, + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("ra1_des", right_schema)?, + options: SortOptions { + descending: true, + nulls_first: true, + }, + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + ]); + let filter_expr = join_expr_tests_fixture_i32( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 5, + side: JoinSide::Left, + }, + ColumnIndex { + index: 5, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn build_null_columns_first() -> Result<()> { + let join_type = JoinType::Full; + let case_expr = 1; + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table((10, 11), 8)?; + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = [PhysicalSortExpr { + expr: col("l_asc_null_first", left_schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("r_asc_null_first", right_schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + ]); + let filter_expr = join_expr_tests_fixture_i32( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 6, + side: JoinSide::Left, + }, + ColumnIndex { + index: 6, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn build_null_columns_last() -> Result<()> { + let join_type = JoinType::Full; + let case_expr = 1; + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table((10, 11), 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = [PhysicalSortExpr { + expr: col("l_asc_null_last", left_schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("r_asc_null_last", right_schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + ]); + let filter_expr = join_expr_tests_fixture_i32( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 7, + side: JoinSide::Left, + }, + ColumnIndex { + index: 7, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn build_null_columns_first_descending() -> Result<()> { + let join_type = JoinType::Full; + let cardinality = (10, 11); + let case_expr = 1; + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = [PhysicalSortExpr { + expr: col("l_desc_null_first", left_schema)?, + options: SortOptions { + descending: true, + nulls_first: true, + }, + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("r_desc_null_first", right_schema)?, + options: SortOptions { + descending: true, + nulls_first: true, + }, + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Int32, true), + Field::new("right", DataType::Int32, true), + ]); + let filter_expr = join_expr_tests_fixture_i32( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 8, + side: JoinSide::Left, + }, + ColumnIndex { + index: 8, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn complex_join_all_one_ascending_numeric_missing_stat() -> Result<()> { + let cardinality = (3, 4); + let join_type = JoinType::Full; + + // a + b > c + 10 AND a + b < c + 100 + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = [PhysicalSortExpr { + expr: col("la1", left_schema)?, + options: SortOptions::default(), + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("ra1", right_schema)?, + options: SortOptions::default(), + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + Field::new("2", DataType::Int32, true), + ]); + let filter_expr = complicated_filter(&intermediate_schema)?; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 4, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[tokio::test(flavor = "multi_thread")] + async fn complex_join_all_one_ascending_equivalence() -> Result<()> { + let cardinality = (3, 4); + let join_type = JoinType::Full; + + // a + b > c + 10 AND a + b < c + 100 + let config = SessionConfig::new().with_repartition_joins(false); + // let session_ctx = SessionContext::with_config(config); + // let task_ctx = session_ctx.task_ctx(); + let task_ctx = Arc::new(TaskContext::default().with_session_config(config)); + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = vec![ + [PhysicalSortExpr { + expr: col("la1", left_schema)?, + options: SortOptions::default(), + }] + .into(), + [PhysicalSortExpr { + expr: col("la2", left_schema)?, + options: SortOptions::default(), + }] + .into(), + ]; + + let right_sorted = [PhysicalSortExpr { + expr: col("ra1", right_schema)?, + options: SortOptions::default(), + }] + .into(); + + let (left, right) = create_memory_table( + left_partition, + right_partition, + left_sorted, + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("0", DataType::Int32, true), + Field::new("1", DataType::Int32, true), + Field::new("2", DataType::Int32, true), + ]); + let filter_expr = complicated_filter(&intermediate_schema)?; + let column_indices = vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 4, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn testing_with_temporal_columns( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values( + (4, 5), + (12, 17), + )] + cardinality: (i32, i32), + #[values(0, 1, 2)] case_expr: usize, + ) -> Result<()> { + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + let left_sorted = [PhysicalSortExpr { + expr: col("lt1", left_schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("rt1", right_schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + let intermediate_schema = Schema::new(vec![ + Field::new( + "left", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new( + "right", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + ]); + let filter_expr = join_expr_tests_fixture_temporal( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + &intermediate_schema, + )?; + let column_indices = vec![ + ColumnIndex { + index: 3, + side: JoinSide::Left, + }, + ColumnIndex { + index: 3, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn test_with_interval_columns( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values( + (4, 5), + (12, 17), + )] + cardinality: (i32, i32), + ) -> Result<()> { + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + let left_sorted = [PhysicalSortExpr { + expr: col("li1", left_schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("ri1", right_schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Interval(IntervalUnit::DayTime), false), + Field::new("right", DataType::Interval(IntervalUnit::DayTime), false), + ]); + let filter_expr = join_expr_tests_fixture_temporal( + 0, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + &intermediate_schema, + )?; + let column_indices = vec![ + ColumnIndex { + index: 9, + side: JoinSide::Left, + }, + ColumnIndex { + index: 9, + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + + Ok(()) + } + + #[rstest] + #[tokio::test(flavor = "multi_thread")] + async fn testing_ascending_float_pruning( + #[values( + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::RightSemi, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::LeftMark, + JoinType::RightAnti, + JoinType::RightMark, + JoinType::Full + )] + join_type: JoinType, + #[values( + (4, 5), + (12, 17), + )] + cardinality: (i32, i32), + #[values(0, 1, 2, 3, 4, 5)] case_expr: usize, + ) -> Result<()> { + let session_config = SessionConfig::new().with_repartition_joins(false); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + let (left_partition, right_partition) = get_or_create_table(cardinality, 8)?; + + let left_schema = &left_partition[0].schema(); + let right_schema = &right_partition[0].schema(); + let left_sorted = [PhysicalSortExpr { + expr: col("l_float", left_schema)?, + options: SortOptions::default(), + }] + .into(); + let right_sorted = [PhysicalSortExpr { + expr: col("r_float", right_schema)?, + options: SortOptions::default(), + }] + .into(); + let (left, right) = create_memory_table( + left_partition, + right_partition, + vec![left_sorted], + vec![right_sorted], + )?; + + let on = vec![(col("lc1", left_schema)?, col("rc1", right_schema)?)]; + + let intermediate_schema = Schema::new(vec![ + Field::new("left", DataType::Float64, true), + Field::new("right", DataType::Float64, true), + ]); + let filter_expr = join_expr_tests_fixture_f64( + case_expr, + col("left", &intermediate_schema)?, + col("right", &intermediate_schema)?, + ); + let column_indices = vec![ + ColumnIndex { + index: 10, // l_float + side: JoinSide::Left, + }, + ColumnIndex { + index: 10, // r_float + side: JoinSide::Right, + }, + ]; + let filter = + JoinFilter::new(filter_expr, column_indices, Arc::new(intermediate_schema)); + + experiment(left, right, Some(filter), join_type, on, task_ctx).await?; + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/test_utils.rs b/native/vendor/datafusion-physical-plan/src/joins/test_utils.rs new file mode 100644 index 0000000..0455fb2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/test_utils.rs
@@ -0,0 +1,613 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This file has test utils for hash joins + +use std::sync::Arc; + +use crate::joins::utils::{JoinFilter, JoinOn}; +use crate::joins::{ + HashJoinExec, PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, +}; +use crate::repartition::RepartitionExec; +use crate::test::TestMemoryExec; +use crate::{ExecutionPlan, ExecutionPlanProperties, Partitioning, common}; + +use arrow::array::{ + ArrayRef, Float64Array, Int32Array, IntervalDayTimeArray, RecordBatch, + TimestampMillisecondArray, types::IntervalDayTime, +}; +use arrow::datatypes::{DataType, Schema}; +use arrow::util::pretty::pretty_format_batches; +use datafusion_common::{NullEquality, Result, ScalarValue}; +use datafusion_execution::TaskContext; +use datafusion_expr::{JoinType, Operator}; +use datafusion_physical_expr::expressions::{binary, cast, col, lit}; +use datafusion_physical_expr::intervals::test_utils::{ + gen_conjunctive_numerical_expr, gen_conjunctive_temporal_expr, +}; +use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; + +use rand::prelude::StdRng; +use rand::{Rng, SeedableRng}; + +pub fn compare_batches(collected_1: &[RecordBatch], collected_2: &[RecordBatch]) { + let left_row_num: usize = collected_1.iter().map(|batch| batch.num_rows()).sum(); + let right_row_num: usize = collected_2.iter().map(|batch| batch.num_rows()).sum(); + if left_row_num == 0 && right_row_num == 0 { + return; + } + // compare + let first_formatted = pretty_format_batches(collected_1).unwrap().to_string(); + let second_formatted = pretty_format_batches(collected_2).unwrap().to_string(); + + let mut first_lines: Vec<&str> = first_formatted.trim().lines().collect(); + first_lines.sort_unstable(); + + let mut second_lines: Vec<&str> = second_formatted.trim().lines().collect(); + second_lines.sort_unstable(); + + for (i, (first_line, second_line)) in + first_lines.iter().zip(&second_lines).enumerate() + { + assert_eq!((i, first_line), (i, second_line)); + } +} + +pub async fn partitioned_sym_join_with_filter( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: Option<JoinFilter>, + join_type: &JoinType, + null_equality: NullEquality, + context: Arc<TaskContext>, +) -> Result<Vec<RecordBatch>> { + let partition_count = 4; + + let left_expr = on + .iter() + .map(|(l, _)| Arc::clone(l) as _) + .collect::<Vec<_>>(); + + let right_expr = on + .iter() + .map(|(_, r)| Arc::clone(r) as _) + .collect::<Vec<_>>(); + + let join = SymmetricHashJoinExec::try_new( + Arc::new(RepartitionExec::try_new( + Arc::clone(&left), + Partitioning::Hash(left_expr, partition_count), + )?), + Arc::new(RepartitionExec::try_new( + Arc::clone(&right), + Partitioning::Hash(right_expr, partition_count), + )?), + on, + filter, + join_type, + null_equality, + left.output_ordering().cloned(), + right.output_ordering().cloned(), + StreamJoinPartitionMode::Partitioned, + )?; + + let mut batches = vec![]; + for i in 0..partition_count { + let stream = join.execute(i, Arc::clone(&context))?; + let more_batches = common::collect(stream).await?; + batches.extend( + more_batches + .into_iter() + .filter(|b| b.num_rows() > 0) + .collect::<Vec<_>>(), + ); + } + + Ok(batches) +} + +pub async fn partitioned_hash_join_with_filter( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + on: JoinOn, + filter: Option<JoinFilter>, + join_type: &JoinType, + null_equality: NullEquality, + context: Arc<TaskContext>, +) -> Result<Vec<RecordBatch>> { + let partition_count = 4; + let (left_expr, right_expr) = on + .iter() + .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) + .unzip(); + + let join = Arc::new(HashJoinExec::try_new( + Arc::new(RepartitionExec::try_new( + left, + Partitioning::Hash(left_expr, partition_count), + )?), + Arc::new(RepartitionExec::try_new( + right, + Partitioning::Hash(right_expr, partition_count), + )?), + on, + filter, + join_type, + None, + PartitionMode::Partitioned, + null_equality, + false, // null_aware + )?); + + let mut batches = vec![]; + for i in 0..partition_count { + let stream = join.execute(i, Arc::clone(&context))?; + let more_batches = common::collect(stream).await?; + batches.extend( + more_batches + .into_iter() + .filter(|b| b.num_rows() > 0) + .collect::<Vec<_>>(), + ); + } + + Ok(batches) +} + +pub fn split_record_batches( + batch: &RecordBatch, + batch_size: usize, +) -> Result<Vec<RecordBatch>> { + let row_num = batch.num_rows(); + let number_of_batch = row_num / batch_size; + let mut sizes = vec![batch_size; number_of_batch]; + sizes.push(row_num - (batch_size * number_of_batch)); + let mut result = vec![]; + for (i, size) in sizes.iter().enumerate() { + result.push(batch.slice(i * batch_size, *size)); + } + Ok(result) +} + +struct AscendingRandomFloatIterator { + prev: f64, + max: f64, + rng: StdRng, +} + +impl AscendingRandomFloatIterator { + fn new(min: f64, max: f64) -> Self { + let mut rng = StdRng::seed_from_u64(42); + let initial = rng.random_range(min..max); + AscendingRandomFloatIterator { + prev: initial, + max, + rng, + } + } +} + +impl Iterator for AscendingRandomFloatIterator { + type Item = f64; + + fn next(&mut self) -> Option<Self::Item> { + let value = self.rng.random_range(self.prev..self.max); + self.prev = value; + Some(value) + } +} + +pub fn join_expr_tests_fixture_temporal( + expr_id: usize, + left_col: Arc<dyn PhysicalExpr>, + right_col: Arc<dyn PhysicalExpr>, + schema: &Schema, +) -> Result<Arc<dyn PhysicalExpr>> { + match expr_id { + // constructs ((left_col - INTERVAL '100ms') > (right_col - INTERVAL '200ms')) AND ((left_col - INTERVAL '450ms') < (right_col - INTERVAL '300ms')) + 0 => gen_conjunctive_temporal_expr( + left_col, + right_col, + Operator::Minus, + Operator::Minus, + Operator::Minus, + Operator::Minus, + ScalarValue::new_interval_dt(0, 100), // 100 ms + ScalarValue::new_interval_dt(0, 200), // 200 ms + ScalarValue::new_interval_dt(0, 450), // 450 ms + ScalarValue::new_interval_dt(0, 300), // 300 ms + schema, + ), + // constructs ((left_col - TIMESTAMP '2023-01-01:12.00.03') > (right_col - TIMESTAMP '2023-01-01:12.00.01')) AND ((left_col - TIMESTAMP '2023-01-01:12.00.00') < (right_col - TIMESTAMP '2023-01-01:12.00.02')) + 1 => gen_conjunctive_temporal_expr( + left_col, + right_col, + Operator::Minus, + Operator::Minus, + Operator::Minus, + Operator::Minus, + ScalarValue::TimestampMillisecond(Some(1672574403000), None), // 2023-01-01:12.00.03 + ScalarValue::TimestampMillisecond(Some(1672574401000), None), // 2023-01-01:12.00.01 + ScalarValue::TimestampMillisecond(Some(1672574400000), None), // 2023-01-01:12.00.00 + ScalarValue::TimestampMillisecond(Some(1672574402000), None), // 2023-01-01:12.00.02 + schema, + ), + // constructs ((left_col - DURATION '3 secs') > (right_col - DURATION '2 secs')) AND ((left_col - DURATION '5 secs') < (right_col - DURATION '4 secs')) + 2 => gen_conjunctive_temporal_expr( + left_col, + right_col, + Operator::Minus, + Operator::Minus, + Operator::Minus, + Operator::Minus, + ScalarValue::DurationMillisecond(Some(3000)), // 3 secs + ScalarValue::DurationMillisecond(Some(2000)), // 2 secs + ScalarValue::DurationMillisecond(Some(5000)), // 5 secs + ScalarValue::DurationMillisecond(Some(4000)), // 4 secs + schema, + ), + _ => unreachable!(), + } +} + +// It creates join filters for different type of fields for testing. +macro_rules! join_expr_tests { + ($func_name:ident, $type:ty, $SCALAR:ident) => { + pub fn $func_name( + expr_id: usize, + left_col: Arc<dyn PhysicalExpr>, + right_col: Arc<dyn PhysicalExpr>, + ) -> Arc<dyn PhysicalExpr> { + match expr_id { + // left_col + 1 > right_col + 5 AND left_col + 3 < right_col + 10 + 0 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Plus, + Operator::Plus, + Operator::Plus, + Operator::Plus, + ), + ScalarValue::$SCALAR(Some(1 as $type)), + ScalarValue::$SCALAR(Some(5 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + ScalarValue::$SCALAR(Some(10 as $type)), + (Operator::Gt, Operator::Lt), + ), + // left_col - 1 > right_col + 3 AND left_col + 3 < right_col + 15 + 1 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Minus, + Operator::Plus, + Operator::Plus, + Operator::Plus, + ), + ScalarValue::$SCALAR(Some(1 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + ScalarValue::$SCALAR(Some(15 as $type)), + (Operator::Gt, Operator::Lt), + ), + // left_col - 1 > right_col + 5 AND left_col - 3 < right_col + 10 + 2 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Minus, + Operator::Plus, + Operator::Minus, + Operator::Plus, + ), + ScalarValue::$SCALAR(Some(1 as $type)), + ScalarValue::$SCALAR(Some(5 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + ScalarValue::$SCALAR(Some(10 as $type)), + (Operator::Gt, Operator::Lt), + ), + // left_col - 10 > right_col - 5 AND left_col - 3 < right_col + 10 + 3 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Minus, + Operator::Minus, + Operator::Minus, + Operator::Plus, + ), + ScalarValue::$SCALAR(Some(10 as $type)), + ScalarValue::$SCALAR(Some(5 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + ScalarValue::$SCALAR(Some(10 as $type)), + (Operator::Gt, Operator::Lt), + ), + // left_col - 10 > right_col - 5 AND left_col - 30 < right_col - 3 + 4 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Minus, + Operator::Minus, + Operator::Minus, + Operator::Minus, + ), + ScalarValue::$SCALAR(Some(10 as $type)), + ScalarValue::$SCALAR(Some(5 as $type)), + ScalarValue::$SCALAR(Some(30 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + (Operator::Gt, Operator::Lt), + ), + // left_col - 2 >= right_col + 5 AND left_col + 7 <= right_col - 3 + // (filters all input rows) + 5 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Minus, + Operator::Plus, + Operator::Plus, + Operator::Minus, + ), + ScalarValue::$SCALAR(Some(2 as $type)), + ScalarValue::$SCALAR(Some(5 as $type)), + ScalarValue::$SCALAR(Some(7 as $type)), + ScalarValue::$SCALAR(Some(3 as $type)), + (Operator::GtEq, Operator::LtEq), + ), + // left_col + 28 >= right_col - 11 AND left_col + 21 <= right_col + 39 + 6 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Plus, + Operator::Minus, + Operator::Plus, + Operator::Plus, + ), + ScalarValue::$SCALAR(Some(28 as $type)), + ScalarValue::$SCALAR(Some(11 as $type)), + ScalarValue::$SCALAR(Some(21 as $type)), + ScalarValue::$SCALAR(Some(39 as $type)), + (Operator::Gt, Operator::LtEq), + ), + // left_col + 28 >= right_col - 11 AND left_col - 21 <= right_col + 39 + 7 => gen_conjunctive_numerical_expr( + left_col, + right_col, + ( + Operator::Plus, + Operator::Minus, + Operator::Minus, + Operator::Plus, + ), + ScalarValue::$SCALAR(Some(28 as $type)), + ScalarValue::$SCALAR(Some(11 as $type)), + ScalarValue::$SCALAR(Some(21 as $type)), + ScalarValue::$SCALAR(Some(39 as $type)), + (Operator::GtEq, Operator::Lt), + ), + _ => panic!("No case"), + } + } + }; +} + +join_expr_tests!(join_expr_tests_fixture_i32, i32, Int32); +join_expr_tests!(join_expr_tests_fixture_f64, f64, Float64); + +pub fn build_sides_record_batches( + table_size: i32, + key_cardinality: (i32, i32), +) -> Result<(RecordBatch, RecordBatch)> { + let null_ratio: f64 = 0.4; + let duplicate_ratio = 0.4; + let initial_range = 0..table_size; + let index = (table_size as f64 * null_ratio).round() as i32; + let rest_of = index..table_size; + let ordered: ArrayRef = Arc::new(Int32Array::from_iter( + initial_range.clone().collect::<Vec<i32>>(), + )); + let random_ordered = generate_ordered_array(table_size, duplicate_ratio); + let ordered_des = Arc::new(Int32Array::from_iter( + initial_range.clone().rev().collect::<Vec<i32>>(), + )); + let cardinality = Arc::new(Int32Array::from_iter( + initial_range.clone().map(|x| x % 4).collect::<Vec<i32>>(), + )); + let cardinality_key_left = Arc::new(Int32Array::from_iter( + initial_range + .clone() + .map(|x| x % key_cardinality.0) + .collect::<Vec<i32>>(), + )); + let cardinality_key_right = Arc::new(Int32Array::from_iter( + initial_range + .clone() + .map(|x| x % key_cardinality.1) + .collect::<Vec<i32>>(), + )); + let ordered_asc_null_first = Arc::new(Int32Array::from_iter({ + std::iter::repeat_n(None, index as usize) + .chain(rest_of.clone().map(Some)) + .collect::<Vec<Option<i32>>>() + })); + let ordered_asc_null_last = Arc::new(Int32Array::from_iter({ + rest_of + .clone() + .map(Some) + .chain(std::iter::repeat_n(None, index as usize)) + .collect::<Vec<Option<i32>>>() + })); + + let ordered_desc_null_first = Arc::new(Int32Array::from_iter({ + std::iter::repeat_n(None, index as usize) + .chain(rest_of.rev().map(Some)) + .collect::<Vec<Option<i32>>>() + })); + + let time = Arc::new(TimestampMillisecondArray::from( + initial_range + .clone() + .map(|x| x as i64 + 1672531200000) // x + 2023-01-01:00.00.00 + .collect::<Vec<i64>>(), + )); + let interval_time: ArrayRef = Arc::new(IntervalDayTimeArray::from( + initial_range + .map(|x| IntervalDayTime { + days: 0, + milliseconds: x * 100, + }) // x * 100ms + .collect::<Vec<_>>(), + )); + + let float_asc = Arc::new(Float64Array::from_iter_values( + AscendingRandomFloatIterator::new(0., table_size as f64) + .take(table_size as usize), + )); + + let left = RecordBatch::try_from_iter(vec![ + ("la1", Arc::clone(&ordered)), + ("lb1", Arc::clone(&cardinality) as ArrayRef), + ("lc1", cardinality_key_left), + ("lt1", Arc::clone(&time) as ArrayRef), + ("la2", Arc::clone(&ordered)), + ("la1_des", Arc::clone(&ordered_des) as ArrayRef), + ( + "l_asc_null_first", + Arc::clone(&ordered_asc_null_first) as ArrayRef, + ), + ( + "l_asc_null_last", + Arc::clone(&ordered_asc_null_last) as ArrayRef, + ), + ( + "l_desc_null_first", + Arc::clone(&ordered_desc_null_first) as ArrayRef, + ), + ("li1", Arc::clone(&interval_time)), + ("l_float", Arc::clone(&float_asc) as ArrayRef), + ("l_random_ordered", Arc::clone(&random_ordered) as ArrayRef), + ])?; + let right = RecordBatch::try_from_iter(vec![ + ("ra1", Arc::clone(&ordered)), + ("rb1", cardinality), + ("rc1", cardinality_key_right), + ("rt1", time), + ("ra2", ordered), + ("ra1_des", ordered_des), + ("r_asc_null_first", ordered_asc_null_first), + ("r_asc_null_last", ordered_asc_null_last), + ("r_desc_null_first", ordered_desc_null_first), + ("ri1", interval_time), + ("r_float", float_asc), + ("r_random_ordered", random_ordered), + ])?; + Ok((left, right)) +} + +pub fn create_memory_table( + left_partition: Vec<RecordBatch>, + right_partition: Vec<RecordBatch>, + left_sorted: Vec<LexOrdering>, + right_sorted: Vec<LexOrdering>, +) -> Result<(Arc<dyn ExecutionPlan>, Arc<dyn ExecutionPlan>)> { + let left_schema = left_partition[0].schema(); + let left = TestMemoryExec::try_new(&[left_partition], left_schema, None)? + .try_with_sort_information(left_sorted)?; + let right_schema = right_partition[0].schema(); + let right = TestMemoryExec::try_new(&[right_partition], right_schema, None)? + .try_with_sort_information(right_sorted)?; + let left = Arc::new(left); + let right = Arc::new(right); + Ok(( + Arc::new(TestMemoryExec::update_cache(&left)), + Arc::new(TestMemoryExec::update_cache(&right)), + )) +} + +/// Filter expr for a + b > c + 10 AND a + b < c + 100 +pub(crate) fn complicated_filter( + filter_schema: &Schema, +) -> Result<Arc<dyn PhysicalExpr>> { + let left_expr = binary( + cast( + binary( + col("0", filter_schema)?, + Operator::Plus, + col("1", filter_schema)?, + filter_schema, + )?, + filter_schema, + DataType::Int64, + )?, + Operator::Gt, + binary( + cast(col("2", filter_schema)?, filter_schema, DataType::Int64)?, + Operator::Plus, + lit(ScalarValue::Int64(Some(10))), + filter_schema, + )?, + filter_schema, + )?; + + let right_expr = binary( + cast( + binary( + col("0", filter_schema)?, + Operator::Plus, + col("1", filter_schema)?, + filter_schema, + )?, + filter_schema, + DataType::Int64, + )?, + Operator::Lt, + binary( + cast(col("2", filter_schema)?, filter_schema, DataType::Int64)?, + Operator::Plus, + lit(ScalarValue::Int64(Some(100))), + filter_schema, + )?, + filter_schema, + )?; + binary(left_expr, Operator::And, right_expr, filter_schema) +} + +fn generate_ordered_array(size: i32, duplicate_ratio: f32) -> Arc<Int32Array> { + let mut rng = StdRng::seed_from_u64(42); + let unique_count = (size as f32 * (1.0 - duplicate_ratio)) as i32; + + // Generate unique random values + let mut values: Vec<i32> = (0..unique_count) + .map(|_| rng.random_range(1..500)) // Modify as per your range + .collect(); + + // Duplicate the values according to the duplicate ratio + for _ in 0..(size - unique_count) { + let index = rng.random_range(0..unique_count); + values.push(values[index as usize]); + } + + // Sort the values to ensure they are ordered + values.sort(); + + Arc::new(Int32Array::from_iter(values)) +}
diff --git a/native/vendor/datafusion-physical-plan/src/joins/utils.rs b/native/vendor/datafusion-physical-plan/src/joins/utils.rs new file mode 100644 index 0000000..20467a7 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/joins/utils.rs
@@ -0,0 +1,4968 @@ +// 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. + +//! Join related functionality used both on logical and physical plans + +use std::cmp::{Ordering, min}; +use std::collections::HashSet; +use std::fmt::{self, Debug}; +use std::future::Future; +use std::iter::once; +use std::ops::Range; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use crate::joins::SharedBitmapBuilder; +use crate::metrics::{ + self, BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + MetricType, +}; +use crate::projection::{ProjectionExec, ProjectionExpr}; +use crate::{ + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, + RangePartitioning, Statistics, +}; +// compatibility +pub use super::join_filter::JoinFilter; +pub use super::join_hash_map::JoinHashMapType; +pub use crate::joins::{JoinOn, JoinOnRef}; + +use arrow::array::{ + Array, ArrowPrimitiveType, BooleanBufferBuilder, NativeAdapter, PrimitiveArray, + RecordBatch, RecordBatchOptions, UInt32Array, UInt32Builder, UInt64Array, + builder::UInt64Builder, downcast_array, new_null_array, +}; +use arrow::array::{ + ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, Date32Array, Date64Array, + Decimal128Array, FixedSizeBinaryArray, Float32Array, Float64Array, Int8Array, + Int16Array, Int32Array, Int64Array, LargeBinaryArray, LargeStringArray, StringArray, + StringViewArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt16Array, +}; +use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::compute::{self, take}; +use arrow::datatypes::{ + ArrowNativeType, Field, Schema, SchemaBuilder, UInt32Type, UInt64Type, +}; +use arrow_ord::ord::{DynComparator, make_comparator}; +use arrow_schema::{DataType, SortOptions, TimeUnit}; +use datafusion_common::cast::as_boolean_array; +use datafusion_common::hash_utils::RandomState; +use datafusion_common::hash_utils::create_hashes; +use datafusion_common::stats::Precision; +use datafusion_common::utils::normalize_float_zero; +use datafusion_common::{ + DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, + internal_datafusion_err, not_impl_err, plan_err, +}; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{ + LexOrdering, PhysicalExpr, PhysicalExprRef, add_offset_to_expr, + add_offset_to_physical_sort_exprs, +}; + +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::future::{BoxFuture, Shared}; +use futures::{FutureExt, ready}; +use parking_lot::Mutex; + +/// Checks whether the schemas "left" and "right" and columns "on" represent a valid join. +/// They are valid whenever their columns' intersection equals the set `on` +pub fn check_join_is_valid(left: &Schema, right: &Schema, on: JoinOnRef) -> Result<()> { + let left: HashSet<Column> = left + .fields() + .iter() + .enumerate() + .map(|(idx, f)| Column::new(f.name(), idx)) + .collect(); + let right: HashSet<Column> = right + .fields() + .iter() + .enumerate() + .map(|(idx, f)| Column::new(f.name(), idx)) + .collect(); + + check_join_set_is_valid(&left, &right, on) +} + +/// Checks whether the sets left, right and on compose a valid join. +/// They are valid whenever their intersection equals the set `on` +fn check_join_set_is_valid( + left: &HashSet<Column>, + right: &HashSet<Column>, + on: &[(PhysicalExprRef, PhysicalExprRef)], +) -> Result<()> { + let on_left = &on + .iter() + .flat_map(|on| collect_columns(&on.0)) + .collect::<HashSet<_>>(); + let left_missing = on_left.difference(left).collect::<HashSet<_>>(); + + let on_right = &on + .iter() + .flat_map(|on| collect_columns(&on.1)) + .collect::<HashSet<_>>(); + let right_missing = on_right.difference(right).collect::<HashSet<_>>(); + + if !left_missing.is_empty() | !right_missing.is_empty() { + return plan_err!( + "The left or right side of the join does not have all columns on \"on\": \nMissing on the left: {left_missing:?}\nMissing on the right: {right_missing:?}" + ); + }; + + Ok(()) +} + +/// Adjust the right out partitioning to new Column Index +pub fn adjust_right_output_partitioning( + right_partitioning: &Partitioning, + left_columns_len: usize, +) -> Result<Partitioning> { + let result = match right_partitioning { + Partitioning::Hash(exprs, size) => { + let new_exprs = exprs + .iter() + .map(|expr| add_offset_to_expr(Arc::clone(expr), left_columns_len as _)) + .collect::<Result<_>>()?; + Partitioning::Hash(new_exprs, *size) + } + Partitioning::Range(range) => { + let ordering = add_offset_to_physical_sort_exprs( + range.ordering().iter().cloned(), + left_columns_len as _, + )?; + let ordering = LexOrdering::new(ordering).ok_or_else(|| { + internal_datafusion_err!( + "Offsetting range partitioning produced an empty ordering" + ) + })?; + Partitioning::Range(RangePartitioning::new( + ordering, + range.split_points().to_vec(), + )) + } + result => result.clone(), + }; + Ok(result) +} + +/// Calculate the output ordering of a given join operation. +pub fn calculate_join_output_ordering( + left_ordering: Option<&LexOrdering>, + right_ordering: Option<&LexOrdering>, + join_type: JoinType, + left_columns_len: usize, + maintains_input_order: &[bool], + probe_side: Option<JoinSide>, +) -> Result<Option<LexOrdering>> { + match maintains_input_order { + [true, false] => { + // Special case, we can prefix ordering of right side with the ordering of left side. + if join_type == JoinType::Inner + && probe_side == Some(JoinSide::Left) + && let Some(right_ordering) = right_ordering.cloned() + { + let right_offset = add_offset_to_physical_sort_exprs( + right_ordering, + left_columns_len as _, + )?; + return if let Some(left_ordering) = left_ordering { + let mut result = left_ordering.clone(); + result.extend(right_offset); + Ok(Some(result)) + } else { + Ok(LexOrdering::new(right_offset)) + }; + } + Ok(left_ordering.cloned()) + } + [false, true] => { + // Special case, we can prefix ordering of left side with the ordering of right side. + if join_type == JoinType::Inner && probe_side == Some(JoinSide::Right) { + return if let Some(right_ordering) = right_ordering.cloned() { + let mut right_offset = add_offset_to_physical_sort_exprs( + right_ordering, + left_columns_len as _, + )?; + if let Some(left_ordering) = left_ordering { + right_offset.extend(left_ordering.clone()); + } + Ok(LexOrdering::new(right_offset)) + } else { + Ok(left_ordering.cloned()) + }; + } + let Some(right_ordering) = right_ordering else { + return Ok(None); + }; + match join_type { + JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => { + add_offset_to_physical_sort_exprs( + right_ordering.clone(), + left_columns_len as _, + ) + .map(LexOrdering::new) + } + _ => Ok(Some(right_ordering.clone())), + } + } + // Doesn't maintain ordering, output ordering is None. + [false, false] => Ok(None), + [true, true] => unreachable!("Cannot maintain ordering of both sides"), + _ => unreachable!("Join operators can not have more than two children"), + } +} + +/// Information about the index and placement (left or right) of the columns +#[derive(Debug, Clone, PartialEq)] +pub struct ColumnIndex { + /// Index of the column + pub index: usize, + /// Whether the column is at the left or right side + pub side: JoinSide, +} + +/// Returns the output field given the input field. Outer joins may +/// insert nulls even if the input was not null +fn output_join_field(old_field: &Field, join_type: &JoinType, is_left: bool) -> Field { + let force_nullable = match join_type { + JoinType::Inner => false, + JoinType::Left => !is_left, // right input is padded with nulls + JoinType::Right => is_left, // left input is padded with nulls + JoinType::Full => true, // both inputs can be padded with nulls + JoinType::LeftSemi => false, // doesn't introduce nulls + JoinType::RightSemi => false, // doesn't introduce nulls + JoinType::LeftAnti => false, // doesn't introduce nulls (or can it??) + JoinType::RightAnti => false, // doesn't introduce nulls (or can it??) + JoinType::LeftMark => false, + JoinType::RightMark => false, + }; + + if force_nullable { + old_field.clone().with_nullable(true) + } else { + old_field.clone() + } +} + +/// Creates a schema for a join operation. +/// The fields from the left side are first +pub fn build_join_schema( + left: &Schema, + right: &Schema, + join_type: &JoinType, +) -> (Schema, Vec<ColumnIndex>) { + let left_fields = || { + left.fields() + .iter() + .map(|f| output_join_field(f, join_type, true)) + .enumerate() + .map(|(index, f)| { + ( + f, + ColumnIndex { + index, + side: JoinSide::Left, + }, + ) + }) + }; + + let right_fields = || { + right + .fields() + .iter() + .map(|f| output_join_field(f, join_type, false)) + .enumerate() + .map(|(index, f)| { + ( + f, + ColumnIndex { + index, + side: JoinSide::Right, + }, + ) + }) + }; + + let (fields, column_indices): (SchemaBuilder, Vec<ColumnIndex>) = match join_type { + JoinType::Inner | JoinType::Left | JoinType::Full | JoinType::Right => { + // left then right + left_fields().chain(right_fields()).unzip() + } + JoinType::LeftSemi | JoinType::LeftAnti => left_fields().unzip(), + JoinType::LeftMark => { + let right_field = once(( + Field::new("mark", DataType::Boolean, false), + ColumnIndex { + index: 0, + side: JoinSide::None, + }, + )); + left_fields().chain(right_field).unzip() + } + JoinType::RightSemi | JoinType::RightAnti => right_fields().unzip(), + JoinType::RightMark => { + let left_field = once(( + Field::new("mark", DataType::Boolean, false), + ColumnIndex { + index: 0, + side: JoinSide::None, + }, + )); + right_fields().chain(left_field).unzip() + } + }; + + let (schema1, schema2) = match join_type { + JoinType::Right + | JoinType::RightSemi + | JoinType::RightAnti + | JoinType::RightMark => (left, right), + _ => (right, left), + }; + + let metadata = schema1 + .metadata() + .clone() + .into_iter() + .chain(schema2.metadata().clone()) + .collect(); + + (fields.finish().with_metadata(metadata), column_indices) +} + +/// A [`OnceAsync`] runs an `async` closure once, where multiple calls to +/// [`OnceAsync::try_once`] return a [`OnceFut`] that resolves to the result of the +/// same computation. +/// +/// This is useful for joins where the results of one child are needed to proceed +/// with multiple output stream +/// +/// +/// For example, in a hash join, one input is buffered and shared across +/// potentially multiple output partitions. Each output partition must wait for +/// the hash table to be built before proceeding. +/// +/// Each output partition waits on the same `OnceAsync` before proceeding. +pub(crate) struct OnceAsync<T> { + fut: Mutex<Option<SharedResult<OnceFut<T>>>>, +} + +impl<T> Default for OnceAsync<T> { + fn default() -> Self { + Self { + fut: Mutex::new(None), + } + } +} + +impl<T> Debug for OnceAsync<T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "OnceAsync") + } +} + +impl<T: 'static> OnceAsync<T> { + /// If this is the first call to this function on this object, will invoke + /// `f` to obtain a future and return a [`OnceFut`] referring to this. `f` + /// may fail, in which case its error is returned. + /// + /// If this is not the first call, will return a [`OnceFut`] referring + /// to the same future as was returned by the first call - or the same + /// error if the initial call to `f` failed. + pub(crate) fn try_once<F, Fut>(&self, f: F) -> Result<OnceFut<T>> + where + F: FnOnce() -> Result<Fut>, + Fut: Future<Output = Result<T>> + Send + 'static, + { + self.fut + .lock() + .get_or_insert_with(|| f().map(OnceFut::new).map_err(Arc::new)) + .clone() + .map_err(DataFusionError::Shared) + } +} + +/// The shared future type used internally within [`OnceAsync`] +type OnceFutPending<T> = Shared<BoxFuture<'static, SharedResult<Arc<T>>>>; + +/// A [`OnceFut`] represents a shared asynchronous computation, that will be evaluated +/// once for all [`Clone`]'s, with [`OnceFut::get`] providing a non-consuming interface +/// to drive the underlying [`Future`] to completion +pub(crate) struct OnceFut<T> { + state: OnceFutState<T>, +} + +impl<T> Clone for OnceFut<T> { + fn clone(&self) -> Self { + Self { + state: self.state.clone(), + } + } +} + +/// A shared state between statistic aggregators for a join +/// operation. +#[derive(Clone, Debug, Default)] +struct PartialJoinStatistics { + pub num_rows: usize, + pub total_byte_size: Precision<usize>, + pub column_statistics: Vec<ColumnStatistics>, +} + +/// Estimates the output statistics for a join operation based on input statistics. +/// +/// # Statistics Propagation +/// +/// This function estimates join output statistics using the following approach: +/// - **Row count estimation**: Uses the `on` parameter (equijoin keys) to estimate +/// output cardinality via [`estimate_join_cardinality`]. The estimation is based on +/// column-level statistics (distinct counts, min/max values) of the join keys. +/// - **Column statistics**: Combines column statistics from both inputs. For join types +/// that preserve all columns (Inner, Left, Right, Full), statistics from both sides +/// are concatenated. For semi/anti joins, the preserved side's statistics are +/// normalized as subset estimates. +/// - **Byte size**: For semi/anti joins, sums normalized column byte-size estimates +/// when every output column has one. Other join types return `Precision::Absent` +/// because join output size is difficult to estimate without knowing the actual data. +/// +/// # The `on` Parameter +/// +/// The `on` parameter represents equijoin keys (e.g., `t1.id = t2.id`). When `on` is +/// empty (as in NestedLoopJoinExec which handles non-equijoin predicates), the +/// cardinality estimation cannot compute selectivity from join keys, and this function +/// returns unknown statistics (`num_rows: Precision::Absent`). +/// +/// # Limitations +/// +/// - Does not account for selectivity of arbitrary join filter expressions +/// (e.g., `(t1.v1 + t2.v1) % 2 = 0`). Such filters, common in NestedLoopJoinExec, +/// are not factored into the cardinality estimation. +/// - Column statistics for inner/outer joins are simply combined from inputs +/// without adjusting for join selectivity (acknowledged in the code as +/// needing "filter selectivity analysis"). +pub(crate) fn estimate_join_statistics( + left_stats: Statistics, + right_stats: Statistics, + on: &JoinOn, + null_equality: NullEquality, + join_type: &JoinType, + schema: &Schema, +) -> Result<Statistics> { + let join_stats = + estimate_join_cardinality(join_type, left_stats, right_stats, on, null_equality); + let (num_rows, total_byte_size, column_statistics) = match join_stats { + Some(stats) => ( + Precision::Inexact(stats.num_rows), + stats.total_byte_size, + stats.column_statistics, + ), + None => ( + Precision::Absent, + Precision::Absent, + Statistics::unknown_column(schema), + ), + }; + Ok(Statistics { + num_rows, + total_byte_size, + column_statistics, + }) +} + +// Estimate the cardinality for the given join with input statistics. +fn estimate_join_cardinality( + join_type: &JoinType, + left_stats: Statistics, + right_stats: Statistics, + on: &JoinOn, + null_equality: NullEquality, +) -> Option<PartialJoinStatistics> { + let on_column_indices = on + .iter() + .map(|(left, right)| equijoin_column_indices(left, right)) + .collect::<Vec<_>>(); + + let (left_key_stats, right_key_stats) = on_column_indices + .iter() + .map(|indices| match indices { + Some((left_index, right_index)) => ( + left_stats.column_statistics[*left_index].clone(), + right_stats.column_statistics[*right_index].clone(), + ), + None => ( + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ), + }) + .unzip::<_, _, Vec<_>, Vec<_>>(); + + match join_type { + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + let ij_cardinality = estimate_inner_join_cardinality( + Statistics { + num_rows: left_stats.num_rows, + total_byte_size: Precision::Absent, + column_statistics: left_key_stats, + }, + Statistics { + num_rows: right_stats.num_rows, + total_byte_size: Precision::Absent, + column_statistics: right_key_stats, + }, + )?; + + // The cardinality for inner join can also be used to estimate + // the cardinality of left/right/full outer joins as long as it + // it is greater than the minimum cardinality constraints of these + // joins (so that we don't underestimate the cardinality). + let cardinality = match join_type { + JoinType::Inner => ij_cardinality, + JoinType::Left => ij_cardinality.max(&left_stats.num_rows), + JoinType::Right => ij_cardinality.max(&right_stats.num_rows), + JoinType::Full => ij_cardinality + .max(&left_stats.num_rows) + .add(&ij_cardinality.max(&right_stats.num_rows)) + .sub(&ij_cardinality), + _ => unreachable!(), + }; + + Some(PartialJoinStatistics { + num_rows: *cardinality.get_value()?, + total_byte_size: Precision::Absent, + // We don't do anything specific here, just combine the existing + // statistics which might yield subpar results (although it is + // true, esp regarding min/max). For a better estimation, we need + // filter selectivity analysis first. + column_statistics: left_stats + .column_statistics + .into_iter() + .chain(right_stats.column_statistics) + .collect(), + }) + } + + JoinType::LeftSemi + | JoinType::RightSemi + | JoinType::LeftAnti + | JoinType::RightAnti => { + let is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti); + let is_anti = matches!(join_type, JoinType::LeftAnti | JoinType::RightAnti); + + let (outer_stats, inner_stats, outer_key_stats, inner_key_stats) = if is_left + { + (left_stats, right_stats, left_key_stats, right_key_stats) + } else { + (right_stats, left_stats, right_key_stats, left_key_stats) + }; + + let outer_rows = *outer_stats.num_rows.get_value()?; + + let outer_join_key_stats = Statistics { + num_rows: outer_stats.num_rows, + total_byte_size: Precision::Absent, + column_statistics: outer_key_stats.clone(), + }; + let inner_join_key_stats = Statistics { + num_rows: inner_stats.num_rows, + total_byte_size: Precision::Absent, + column_statistics: inner_key_stats.clone(), + }; + + let semi_cardinality = + if estimate_disjoint_inputs(&outer_join_key_stats, &inner_join_key_stats) + .is_some() + { + // If join keys are disjoint, no rows will match + Some(0) + } else { + estimate_semi_join_cardinality( + &outer_stats.num_rows, + &inner_stats.num_rows, + &outer_key_stats, + &inner_key_stats, + null_equality, + ) + }; + + // Semi joins keep the matching rows; anti joins keep the rest. When no + // estimate is available, conservatively assume all outer rows pass. + let cardinality = match (semi_cardinality, is_anti) { + (Some(semi), true) => outer_rows.saturating_sub(semi), + (Some(semi), false) => semi, + (None, _) => outer_rows, + }; + + // The outer side is the one whose columns a semi/anti join emits, so + // its statistics are the ones to normalize into the subset estimate. + let Statistics { + num_rows: preserved_num_rows, + column_statistics: preserved_column_statistics, + .. + } = outer_stats; + let preserved_join_key_indices = on_column_indices + .iter() + .filter_map(|&indices| { + indices.map( + |(left_index, right_index)| { + if is_left { left_index } else { right_index } + }, + ) + }) + .collect::<Vec<_>>(); + let column_statistics = normalize_semi_anti_join_column_statistics( + preserved_column_statistics, + &preserved_num_rows, + cardinality, + &preserved_join_key_indices, + is_anti, + null_equality, + ); + let total_byte_size = + total_byte_size_from_column_statistics(&column_statistics); + Some(PartialJoinStatistics { + num_rows: cardinality, + total_byte_size, + column_statistics, + }) + } + + JoinType::LeftMark => { + let num_rows = *left_stats.num_rows.get_value()?; + let mut column_statistics = left_stats.column_statistics; + column_statistics.push(ColumnStatistics::new_unknown()); + Some(PartialJoinStatistics { + num_rows, + total_byte_size: Precision::Absent, + column_statistics, + }) + } + JoinType::RightMark => { + let num_rows = *right_stats.num_rows.get_value()?; + let mut column_statistics = right_stats.column_statistics; + column_statistics.push(ColumnStatistics::new_unknown()); + Some(PartialJoinStatistics { + num_rows, + total_byte_size: Precision::Absent, + column_statistics, + }) + } + } +} + +fn equijoin_column_indices( + left: &PhysicalExprRef, + right: &PhysicalExprRef, +) -> Option<(usize, usize)> { + Some(( + left.downcast_ref::<Column>()?.index(), + right.downcast_ref::<Column>()?.index(), + )) +} + +/// Adjusts the preserved input's column statistics to describe the subset of +/// rows a semi or anti join emits. Most values become estimates (marked +/// inexact) bounded by the smaller output row count: +/// +/// - `null_count` and `byte_size` are scaled by the output/input row ratio. +/// - `distinct_count` is capped at the number of non-null output rows. +/// - `sum_value` is dropped, since the input sum does not apply to the subset. +/// +/// Join-key columns are the exception for `null_count`: under regular SQL +/// equality, null keys never match, so a semi join keeps none of those rows and +/// an anti join keeps all of them. Under null-equal joins, null keys can match +/// and are treated like the rest of the subset. +fn normalize_semi_anti_join_column_statistics( + column_statistics: Vec<ColumnStatistics>, + input_num_rows: &Precision<usize>, + output_num_rows: usize, + join_key_indices: &[usize], + is_anti: bool, + null_equality: NullEquality, +) -> Vec<ColumnStatistics> { + let input_num_rows = input_num_rows.get_value().copied().unwrap_or(0); + + column_statistics + .into_iter() + .enumerate() + .map(|(idx, stats)| { + let mut stats = stats.to_inexact(); + stats.null_count = if join_key_indices.contains(&idx) { + normalize_semi_anti_join_key_null_count( + stats.null_count, + input_num_rows, + output_num_rows, + is_anti, + null_equality, + ) + } else { + scale_subset_count(stats.null_count, input_num_rows, output_num_rows) + .min(&Precision::Inexact(output_num_rows)) + }; + let max_distinct_count = stats + .null_count + .get_value() + .map(|null_count| output_num_rows.saturating_sub(*null_count)) + .unwrap_or(output_num_rows); + stats.distinct_count = stats + .distinct_count + .min(&Precision::Inexact(max_distinct_count)); + stats.byte_size = + scale_subset_count(stats.byte_size, input_num_rows, output_num_rows); + stats.sum_value = Precision::Absent; + stats + }) + .collect() +} + +fn normalize_semi_anti_join_key_null_count( + null_count: Precision<usize>, + input_num_rows: usize, + output_num_rows: usize, + is_anti: bool, + null_equality: NullEquality, +) -> Precision<usize> { + match (is_anti, null_equality) { + (false, NullEquality::NullEqualsNothing) => Precision::Exact(0), + (true, NullEquality::NullEqualsNothing) => null_count + .to_inexact() + .min(&Precision::Inexact(output_num_rows)), + (_, NullEquality::NullEqualsNull) => { + scale_subset_count(null_count, input_num_rows, output_num_rows) + .min(&Precision::Inexact(output_num_rows)) + } + } +} + +// Scale a column-level count to an estimated row subset. Rounding up keeps a +// small non-zero count from disappearing solely because the subset is small. +fn scale_subset_count( + count: Precision<usize>, + input_num_rows: usize, + output_num_rows: usize, +) -> Precision<usize> { + let scaled = match count { + Precision::Exact(count) | Precision::Inexact(count) => { + if input_num_rows == 0 { + 0 + } else { + (count as u128 * output_num_rows as u128).div_ceil(input_num_rows as u128) + as usize + } + } + Precision::Absent => return Precision::Absent, + }; + + Precision::Inexact(scaled) +} + +fn total_byte_size_from_column_statistics( + column_statistics: &[ColumnStatistics], +) -> Precision<usize> { + column_statistics + .iter() + .map(|stats| stats.byte_size.get_value().copied()) + .try_fold(0usize, |acc, byte_size| { + byte_size.map(|byte_size| acc.saturating_add(byte_size)) + }) + .map(Precision::Inexact) + .unwrap_or(Precision::Absent) +} + +/// Estimate the inner join cardinality by using the basic building blocks of +/// column-level statistics and the total row count. This is a very naive and +/// a very conservative implementation that can quickly give up if there is not +/// enough input statistics. +fn estimate_inner_join_cardinality( + left_stats: Statistics, + right_stats: Statistics, +) -> Option<Precision<usize>> { + // Immediately return if inputs considered as non-overlapping + if let Some(estimation) = estimate_disjoint_inputs(&left_stats, &right_stats) { + return Some(estimation); + }; + + let Statistics { + num_rows: left_num_rows, + column_statistics: left_column_statistics, + .. + } = left_stats; + let Statistics { + num_rows: right_num_rows, + column_statistics: right_column_statistics, + .. + } = right_stats; + + if left_num_rows == Precision::Exact(0) || right_num_rows == Precision::Exact(0) { + return Some(Precision::Exact(0)); + } + if left_num_rows == Precision::Inexact(0) || right_num_rows == Precision::Inexact(0) { + return Some(Precision::Inexact(0)); + } + + // Follow Spark Catalyst's conservative NDV join estimate: for multi-key + // joins, use the most selective key instead of multiplying all key denominators. + let mut join_selectivity = Precision::Absent; + for (left_stat, right_stat) in left_column_statistics + .iter() + .zip(right_column_statistics.iter()) + { + let left_max_distinct = max_distinct_count(&left_num_rows, left_stat); + let right_max_distinct = max_distinct_count(&right_num_rows, right_stat); + let max_distinct = left_max_distinct.max(&right_max_distinct); + if max_distinct.get_value().is_some() { + // Seems like there are a few implementations of this algorithm that implement + // exponential decay for the selectivity (like Hive's Optiq Optimizer). Needs + // further exploration. + join_selectivity = if join_selectivity.get_value().is_some() { + join_selectivity.max(&max_distinct) + } else { + max_distinct + }; + } + } + + // With the assumption that the smaller input's domain is generally represented in the bigger + // input's domain, we can estimate the inner join's cardinality by taking the cartesian product + // of the two inputs and normalizing it by the selectivity factor. + let left_num_rows = *left_stats.num_rows.get_value()?; + let right_num_rows = *right_stats.num_rows.get_value()?; + // Widen before multiplying so the intermediate Cartesian product does not + // overflow when the normalized cardinality is still representable as usize. + let cartesian_product = (left_num_rows as u128) * (right_num_rows as u128); + let normalized_cardinality = + |value: usize| usize::try_from(cartesian_product / value as u128); + match join_selectivity { + Precision::Exact(value) if value > 0 => Some( + normalized_cardinality(value) + .map(Precision::Exact) + .unwrap_or(Precision::Inexact(usize::MAX)), + ), + Precision::Inexact(value) if value > 0 => Some(Precision::Inexact( + normalized_cardinality(value).unwrap_or(usize::MAX), + )), + // Since we don't have any information about the selectivity (which is derived + // from the number of distinct rows information) we can give up here for now. + // And let other passes handle this (otherwise we would need to produce an + // overestimation using just the cartesian product). + _ => None, + } +} + +/// Estimates if inputs are non-overlapping, using input statistics. +/// If inputs are disjoint, returns zero estimation, otherwise returns None +fn estimate_disjoint_inputs( + left_stats: &Statistics, + right_stats: &Statistics, +) -> Option<Precision<usize>> { + for (left_stat, right_stat) in left_stats + .column_statistics + .iter() + .zip(right_stats.column_statistics.iter()) + { + // If there is no overlap in any of the join columns, this means the join + // itself is disjoint and the cardinality is 0. Though we can only assume + // this when the statistics are exact (since it is a very strong assumption). + let left_min_val = left_stat.min_value.get_value(); + let right_max_val = right_stat.max_value.get_value(); + if left_min_val.is_some() + && right_max_val.is_some() + && left_min_val > right_max_val + { + return Some( + if left_stat.min_value.is_exact().unwrap_or(false) + && right_stat.max_value.is_exact().unwrap_or(false) + { + Precision::Exact(0) + } else { + Precision::Inexact(0) + }, + ); + } + + let left_max_val = left_stat.max_value.get_value(); + let right_min_val = right_stat.min_value.get_value(); + if left_max_val.is_some() + && right_min_val.is_some() + && left_max_val < right_min_val + { + return Some( + if left_stat.max_value.is_exact().unwrap_or(false) + && right_stat.min_value.is_exact().unwrap_or(false) + { + Precision::Exact(0) + } else { + Precision::Inexact(0) + }, + ); + } + } + + None +} + +/// Estimates the number of outer rows that have at least one matching +/// key on the inner side (i.e. semi join cardinality) using NDV +/// (Number of Distinct Values) statistics. +/// +/// Assuming the smaller domain is contained in the larger, the number +/// of overlapping distinct values is `min(outer_ndv, inner_ndv)`. +/// Under the uniformity assumption (each distinct value contributes +/// equally to row counts), the surviving fraction of outer rows is: +/// +/// Under regular SQL equality, null rows cannot match, so each column's +/// selectivity is further reduced by the outer null fraction: +/// +/// ```text +/// null_frac_i = outer_null_count_i / outer_rows +/// selectivity_i = min(outer_ndv_i, inner_ndv_i) / outer_ndv_i * (1 - null_frac_i) +/// ``` +/// +/// For multi-column join keys the overall selectivity is the product +/// of per-column factors: +/// +/// ```text +/// semi_cardinality = outer_rows * product_i(selectivity_i) +/// ``` +/// +/// Anti join cardinality is derived as the complement: +/// `outer_rows - semi_cardinality`. +/// +/// With `NullEqualsNothing`, boundary cases are: +/// * `inner_ndv >= outer_ndv` → selectivity = `1.0 - null_frac` +/// * `null_frac = 1.0` → selectivity = 0.0 (no non-null rows can match) +/// * Missing NDV statistics → returns `None` (fallback to `outer_rows`) +/// +/// PostgreSQL uses a similar approach in `eqjoinsel_semi` +/// (`src/backend/utils/adt/selfuncs.c`). When NDV statistics are +/// available on both sides it computes selectivity as `nd2 / nd1`, +/// which is equivalent to `min(outer_ndv, inner_ndv) / outer_ndv`. +/// If either side lacks statistics it falls back to a default. +fn estimate_semi_join_cardinality( + outer_num_rows: &Precision<usize>, + inner_num_rows: &Precision<usize>, + outer_key_stats: &[ColumnStatistics], + inner_key_stats: &[ColumnStatistics], + null_equality: NullEquality, +) -> Option<usize> { + let outer_rows = *outer_num_rows.get_value()?; + if outer_rows == 0 { + return Some(0); + } + let inner_rows = *inner_num_rows.get_value()?; + if inner_rows == 0 { + return Some(0); + } + + let mut selectivity = 1.0_f64; + let mut has_selectivity_estimate = false; + + for (outer_stat, inner_stat) in outer_key_stats.iter().zip(inner_key_stats.iter()) { + let outer_has_stats = outer_stat.distinct_count.get_value().is_some() + || (outer_stat.min_value.get_value().is_some() + && outer_stat.max_value.get_value().is_some()); + let inner_has_stats = inner_stat.distinct_count.get_value().is_some() + || (inner_stat.min_value.get_value().is_some() + && inner_stat.max_value.get_value().is_some()); + if !outer_has_stats || !inner_has_stats { + continue; + } + + let outer_ndv = max_distinct_count(outer_num_rows, outer_stat); + let inner_ndv = max_distinct_count(inner_num_rows, inner_stat); + + if let (Some(&o), Some(&i)) = (outer_ndv.get_value(), inner_ndv.get_value()) + && o > 0 + { + let null_frac = if null_equality == NullEquality::NullEqualsNothing { + outer_stat + .null_count + .get_value() + .map(|&nc| { + if nc > outer_rows { + 0.0 + } else { + nc as f64 / outer_rows as f64 + } + }) + .unwrap_or(0.0) + } else { + 0.0 + }; + selectivity *= (o.min(i) as f64) / (o as f64) * (1.0 - null_frac); + has_selectivity_estimate = true; + } + } + + if has_selectivity_estimate { + Some((outer_rows as f64 * selectivity).ceil() as usize) + } else { + None + } +} + +/// Estimate the number of maximum distinct values that can be present in the +/// given column from its statistics. If distinct_count is available, uses it +/// directly. Otherwise, if the column is numeric and has min/max values, it +/// estimates the maximum distinct count from those. Otherwise, the num_rows +/// is used. +fn max_distinct_count( + num_rows: &Precision<usize>, + stats: &ColumnStatistics, +) -> Precision<usize> { + match &stats.distinct_count { + &dc @ (Precision::Exact(_) | Precision::Inexact(_)) => { + // NDV can never exceed the number of rows + match num_rows { + Precision::Absent => dc, + _ => { + if dc.get_value() <= num_rows.get_value() { + dc + } else { + num_rows.to_inexact() + } + } + } + } + _ => { + // The number can never be greater than the number of rows we have + // minus the nulls (since they don't count as distinct values). + let result = match num_rows { + Precision::Absent => Precision::Absent, + Precision::Inexact(count) => { + // To safeguard against inexact number of rows (e.g. 0) being smaller than + // an exact null count we need to do a checked subtraction. + match count.checked_sub(*stats.null_count.get_value().unwrap_or(&0)) { + None => Precision::Inexact(0), + Some(non_null_count) => Precision::Inexact(non_null_count), + } + } + Precision::Exact(count) => { + let null_count = *stats.null_count.get_value().unwrap_or(&0); + let non_null_count = count.checked_sub(null_count).unwrap_or(0); + if stats.null_count.is_exact().unwrap_or(false) { + Precision::Exact(non_null_count) + } else { + Precision::Inexact(non_null_count) + } + } + }; + // Cap the estimate using the number of possible values: + if let (Some(min), Some(max)) = + (stats.min_value.get_value(), stats.max_value.get_value()) + && let Some(range_dc) = Interval::try_new(min.clone(), max.clone()) + .ok() + .and_then(|e| e.cardinality()) + { + let range_dc = range_dc as usize; + // Note that the `unwrap` calls in the below statement are safe. + return if result == Precision::Absent + || &range_dc < result.get_value().unwrap() + { + if stats.min_value.is_exact().unwrap() + && stats.max_value.is_exact().unwrap() + { + Precision::Exact(range_dc) + } else { + Precision::Inexact(range_dc) + } + } else { + result + }; + } + + result + } + } +} + +enum OnceFutState<T> { + Pending(OnceFutPending<T>), + Ready(SharedResult<Arc<T>>), +} + +impl<T> Clone for OnceFutState<T> { + fn clone(&self) -> Self { + match self { + Self::Pending(p) => Self::Pending(p.clone()), + Self::Ready(r) => Self::Ready(r.clone()), + } + } +} + +impl<T: 'static> OnceFut<T> { + /// Create a new [`OnceFut`] from a [`Future`] + pub(crate) fn new<Fut>(fut: Fut) -> Self + where + Fut: Future<Output = Result<T>> + Send + 'static, + { + Self { + state: OnceFutState::Pending( + fut.map(|res| res.map(Arc::new).map_err(Arc::new)) + .boxed() + .shared(), + ), + } + } + + /// Get the result of the computation if it is ready, without consuming it + pub(crate) fn get(&mut self, cx: &mut Context<'_>) -> Poll<Result<&T>> { + if let OnceFutState::Pending(fut) = &mut self.state { + let r = ready!(fut.poll_unpin(cx)); + self.state = OnceFutState::Ready(r); + } + + // Cannot use loop as this would trip up the borrow checker + match &self.state { + OnceFutState::Pending(_) => unreachable!(), + OnceFutState::Ready(r) => Poll::Ready( + r.as_ref() + .map(|r| r.as_ref()) + .map_err(DataFusionError::from), + ), + } + } + + /// Get shared reference to the result of the computation if it is ready, without consuming it + pub(crate) fn get_shared(&mut self, cx: &mut Context<'_>) -> Poll<Result<Arc<T>>> { + if let OnceFutState::Pending(fut) = &mut self.state { + let r = ready!(fut.poll_unpin(cx)); + self.state = OnceFutState::Ready(r); + } + + match &self.state { + OnceFutState::Pending(_) => unreachable!(), + OnceFutState::Ready(r) => { + Poll::Ready(r.clone().map_err(DataFusionError::Shared)) + } + } + } +} + +/// Should we use a bitmap to track each incoming right batch's each row's +/// 'joined' status. +/// +/// For example in right joins, we have to use a bit map to track matched +/// right side rows, and later enter a `EmitRightUnmatched` stage to emit +/// unmatched right rows. +pub(crate) fn need_produce_right_in_final(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::Full + | JoinType::Right + | JoinType::RightAnti + | JoinType::RightMark + | JoinType::RightSemi + ) +} + +/// Some type `join_type` of join need to maintain the matched indices bit map for the left side, and +/// use the bit map to generate the part of result of the join. +/// +/// For example of the `Left` join, in each iteration of right side, can get the matched result, but need +/// to maintain the matched indices bit map to get the unmatched row for the left side. +pub(crate) fn need_produce_result_in_final(join_type: JoinType) -> bool { + matches!( + join_type, + JoinType::Left + | JoinType::LeftAnti + | JoinType::LeftSemi + | JoinType::LeftMark + | JoinType::Full + ) +} + +pub(crate) fn get_final_indices_from_shared_bitmap( + shared_bitmap: &SharedBitmapBuilder, + join_type: JoinType, + piecewise: bool, +) -> (UInt64Array, UInt32Array) { + let bitmap = shared_bitmap.lock(); + get_final_indices_from_bit_map(&bitmap, join_type, piecewise) +} + +/// In the end of join execution, need to use bit map of the matched +/// indices to generate the final left and right indices. +/// +/// For example: +/// +/// 1. left_bit_map: `[true, false, true, true, false]` +/// 2. join_type: `Left` +/// +/// The result is: `([1,4], [null, null])` +pub(crate) fn get_final_indices_from_bit_map( + left_bit_map: &BooleanBufferBuilder, + join_type: JoinType, + // We add a flag for whether this is being passed from the `PiecewiseMergeJoin` + // because the bitmap can be for left + right `JoinType`s + piecewise: bool, +) -> (UInt64Array, UInt32Array) { + let left_size = left_bit_map.len(); + if join_type == JoinType::LeftMark || (join_type == JoinType::RightMark && piecewise) + { + let left_indices = (0..left_size as u64).collect::<UInt64Array>(); + let right_indices = (0..left_size) + .map(|idx| left_bit_map.get_bit(idx).then_some(0)) + .collect::<UInt32Array>(); + return (left_indices, right_indices); + } + let left_indices = if join_type == JoinType::LeftSemi + || (join_type == JoinType::RightSemi && piecewise) + { + (0..left_size) + .filter_map(|idx| (left_bit_map.get_bit(idx)).then_some(idx as u64)) + .collect::<UInt64Array>() + } else { + // just for `Left`, `LeftAnti` and `Full` join + // `LeftAnti`, `Left` and `Full` will produce the unmatched left row finally + (0..left_size) + .filter_map(|idx| (!left_bit_map.get_bit(idx)).then_some(idx as u64)) + .collect::<UInt64Array>() + }; + // right_indices + // all the element in the right side is None + let mut builder = UInt32Builder::with_capacity(left_indices.len()); + builder.append_nulls(left_indices.len()); + let right_indices = builder.finish(); + (left_indices, right_indices) +} + +#[expect(clippy::too_many_arguments)] +pub(crate) fn apply_join_filter_to_indices( + build_input_buffer: &RecordBatch, + probe_batch: &RecordBatch, + build_indices: UInt64Array, + probe_indices: UInt32Array, + filter: &JoinFilter, + build_side: JoinSide, + max_intermediate_size: Option<usize>, + join_type: JoinType, +) -> Result<(UInt64Array, UInt32Array)> { + if build_indices.is_empty() && probe_indices.is_empty() { + return Ok((build_indices, probe_indices)); + }; + + let filter_result = if let Some(max_size) = max_intermediate_size { + let mut filter_results = + Vec::with_capacity(build_indices.len().div_ceil(max_size)); + + for i in (0..build_indices.len()).step_by(max_size) { + let end = min(build_indices.len(), i + max_size); + let len = end - i; + let intermediate_batch = build_batch_from_indices( + filter.schema(), + build_input_buffer, + probe_batch, + &build_indices.slice(i, len), + &probe_indices.slice(i, len), + filter.column_indices(), + build_side, + join_type, + )?; + let filter_result = filter + .expression() + .evaluate(&intermediate_batch)? + .into_array(intermediate_batch.num_rows())?; + filter_results.push(filter_result); + } + + let filter_refs: Vec<&dyn Array> = + filter_results.iter().map(|a| a.as_ref()).collect(); + + compute::concat(&filter_refs)? + } else { + let intermediate_batch = build_batch_from_indices( + filter.schema(), + build_input_buffer, + probe_batch, + &build_indices, + &probe_indices, + filter.column_indices(), + build_side, + join_type, + )?; + + filter + .expression() + .evaluate(&intermediate_batch)? + .into_array(intermediate_batch.num_rows())? + }; + + let mask = as_boolean_array(&filter_result)?; + + let left_filtered = compute::filter(&build_indices, mask)?; + let right_filtered = compute::filter(&probe_indices, mask)?; + Ok(( + downcast_array(left_filtered.as_ref()), + downcast_array(right_filtered.as_ref()), + )) +} + +/// Creates a [RecordBatch] with zero columns but the given row count. +/// Used when a join has an empty projection (e.g. `SELECT count(1) ...`). +fn new_empty_schema_batch(schema: &Schema, row_count: usize) -> Result<RecordBatch> { + let options = RecordBatchOptions::new().with_row_count(Some(row_count)); + Ok(RecordBatch::try_new_with_options( + Arc::new(schema.clone()), + vec![], + &options, + )?) +} + +/// Returns a new [RecordBatch] by combining the `left` and `right` according to `indices`. +/// The resulting batch has [Schema] `schema`. +#[expect(clippy::too_many_arguments)] +pub(crate) fn build_batch_from_indices( + schema: &Schema, + build_input_buffer: &RecordBatch, + probe_batch: &RecordBatch, + build_indices: &UInt64Array, + probe_indices: &UInt32Array, + column_indices: &[ColumnIndex], + build_side: JoinSide, + join_type: JoinType, +) -> Result<RecordBatch> { + if schema.fields().is_empty() { + // For RightAnti and RightSemi joins, after `adjust_indices_by_join_type` + // the build_indices were untouched so only probe_indices hold the actual + // row count. + let row_count = match join_type { + JoinType::RightAnti | JoinType::RightSemi => probe_indices.len(), + _ => build_indices.len(), + }; + return new_empty_schema_batch(schema, row_count); + } + + // build the columns of the new [RecordBatch]: + // 1. pick whether the column is from the left or right + // 2. based on the pick, `take` items from the different RecordBatches + let mut columns: Vec<Arc<dyn Array>> = Vec::with_capacity(schema.fields().len()); + + for column_index in column_indices { + let array = if column_index.side == JoinSide::None { + // For mark joins, the mark column is a true if the indices is not null, otherwise it will be false + Arc::new(compute::is_not_null(probe_indices)?) + } else if column_index.side == build_side { + let array = build_input_buffer.column(column_index.index); + if array.is_empty() || build_indices.null_count() == build_indices.len() { + // Outer join would generate a null index when finding no match at our side. + // Therefore, it's possible we are empty but need to populate an n-length null array, + // where n is the length of the index array. + assert_eq!(build_indices.null_count(), build_indices.len()); + new_null_array(array.data_type(), build_indices.len()) + } else { + take(array.as_ref(), build_indices, None)? + } + } else { + let array = probe_batch.column(column_index.index); + if array.is_empty() || probe_indices.null_count() == probe_indices.len() { + assert_eq!(probe_indices.null_count(), probe_indices.len()); + new_null_array(array.data_type(), probe_indices.len()) + } else { + take(array.as_ref(), probe_indices, None)? + } + }; + + columns.push(array); + } + Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?) +} + +/// Returns a new [RecordBatch] for a probe batch when no probe row can find a +/// match: the build-side map is empty, either because the build side has no +/// rows or because none of its rows has a matchable (non-NULL) join key. +/// The resulting batch has [Schema] `schema`. +pub(crate) fn build_batch_empty_build_side( + schema: &Schema, + build_batch: &RecordBatch, + probe_batch: &RecordBatch, + column_indices: &[ColumnIndex], + join_type: JoinType, +) -> Result<RecordBatch> { + if join_type.empty_build_side_produces_empty_result() { + // These join types only return data if the left side is not empty. + return Ok(RecordBatch::new_empty(Arc::new(schema.clone()))); + } + + // The remaining joins return right-side rows and nulls for the left side. + let num_rows = probe_batch.num_rows(); + if schema.fields().is_empty() { + return new_empty_schema_batch(schema, num_rows); + } + + let columns = column_indices + .iter() + .map(|column_index| match column_index.side { + // left -> null array + JoinSide::Left => new_null_array( + build_batch.column(column_index.index).data_type(), + num_rows, + ), + // right -> respective right array + JoinSide::Right => Arc::clone(probe_batch.column(column_index.index)), + // right mark -> unset boolean array as there are no matches on the left side + JoinSide::None => { + Arc::new(BooleanArray::new(BooleanBuffer::new_unset(num_rows), None)) + } + }) + .collect(); + + Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?) +} + +/// The input is the matched indices for left and right and +/// adjust the indices according to the join type +pub(crate) fn adjust_indices_by_join_type( + left_indices: UInt64Array, + right_indices: UInt32Array, + adjust_range: Range<usize>, + join_type: JoinType, + preserve_order_for_right: bool, +) -> Result<(UInt64Array, UInt32Array)> { + match join_type { + JoinType::Inner => { + // matched + Ok((left_indices, right_indices)) + } + JoinType::Left => { + // matched + Ok((left_indices, right_indices)) + // unmatched left row will be produced in the end of loop, and it has been set in the left visited bitmap + } + JoinType::Right => { + // combine the matched and unmatched right result together + append_right_indices( + left_indices, + right_indices, + adjust_range, + preserve_order_for_right, + ) + } + JoinType::Full => { + append_right_indices(left_indices, right_indices, adjust_range, false) + } + JoinType::RightSemi => { + // need to remove the duplicated record in the right side + let right_indices = get_semi_indices(adjust_range, &right_indices); + // the left_indices will not be used later for the `right semi` join + Ok((left_indices, right_indices)) + } + JoinType::RightAnti => { + // need to remove the duplicated record in the right side + // get the anti index for the right side + let right_indices = get_anti_indices(adjust_range, &right_indices); + // the left_indices will not be used later for the `right anti` join + Ok((left_indices, right_indices)) + } + JoinType::RightMark => { + let right_indices = get_mark_indices(&adjust_range, &right_indices); + let left_indices_vec: Vec<u64> = adjust_range.map(|i| i as u64).collect(); + let left_indices = UInt64Array::from(left_indices_vec); + Ok((left_indices, right_indices)) + } + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + // matched or unmatched left row will be produced in the end of loop + // When visit the right batch, we can output the matched left row and don't need to wait the end of loop + Ok(( + UInt64Array::from_iter_values(vec![]), + UInt32Array::from_iter_values(vec![]), + )) + } + } +} + +/// Appends right indices to left indices based on the specified order mode. +/// +/// The function operates in two modes: +/// 1. If `preserve_order_for_right` is true, probe matched and unmatched indices +/// are inserted in order using the `append_probe_indices_in_order()` method. +/// 2. Otherwise, unmatched probe indices are simply appended after matched ones. +/// +/// # Parameters +/// - `left_indices`: UInt64Array of left indices. +/// - `right_indices`: UInt32Array of right indices. +/// - `adjust_range`: Range to adjust the right indices. +/// - `preserve_order_for_right`: Boolean flag to determine the mode of operation. +/// +/// # Returns +/// A tuple of updated `UInt64Array` and `UInt32Array`. +pub(crate) fn append_right_indices( + left_indices: UInt64Array, + right_indices: UInt32Array, + adjust_range: Range<usize>, + preserve_order_for_right: bool, +) -> Result<(UInt64Array, UInt32Array)> { + if preserve_order_for_right { + Ok(append_probe_indices_in_order( + &left_indices, + &right_indices, + adjust_range, + )) + } else { + let right_unmatched_indices = get_anti_indices(adjust_range, &right_indices); + + if right_unmatched_indices.is_empty() { + Ok((left_indices, right_indices)) + } else { + // `into_builder()` can fail here when there is nothing to be filtered and + // left_indices or right_indices has the same reference to the cached indices. + // In that case, we use a slower alternative. + + // the new left indices: left_indices + null array + let mut new_left_indices_builder = + left_indices.into_builder().unwrap_or_else(|left_indices| { + let mut builder = UInt64Builder::with_capacity( + left_indices.len() + right_unmatched_indices.len(), + ); + debug_assert_eq!( + left_indices.null_count(), + 0, + "expected left indices to have no nulls" + ); + builder.append_slice(left_indices.values()); + builder + }); + new_left_indices_builder.append_nulls(right_unmatched_indices.len()); + let new_left_indices = UInt64Array::from(new_left_indices_builder.finish()); + + // the new right indices: right_indices + right_unmatched_indices + let mut new_right_indices_builder = right_indices + .into_builder() + .unwrap_or_else(|right_indices| { + let mut builder = UInt32Builder::with_capacity( + right_indices.len() + right_unmatched_indices.len(), + ); + debug_assert_eq!( + right_indices.null_count(), + 0, + "expected right indices to have no nulls" + ); + builder.append_slice(right_indices.values()); + builder + }); + debug_assert_eq!( + right_unmatched_indices.null_count(), + 0, + "expected right unmatched indices to have no nulls" + ); + new_right_indices_builder.append_slice(right_unmatched_indices.values()); + let new_right_indices = UInt32Array::from(new_right_indices_builder.finish()); + + Ok((new_left_indices, new_right_indices)) + } + } +} + +/// Returns `range` indices which are not present in `input_indices`. +/// +/// `input_indices` must be sorted ascending and contain no nulls. +pub(crate) fn get_anti_indices<T: ArrowPrimitiveType>( + range: Range<usize>, + input_indices: &PrimitiveArray<T>, +) -> PrimitiveArray<T> +where + NativeAdapter<T>: From<<T as ArrowPrimitiveType>::Native>, +{ + debug_assert_eq!( + input_indices.null_count(), + 0, + "get_anti_indices requires non-null input_indices" + ); + debug_assert!( + input_indices + .values() + .windows(2) + .all(|w| w[0].as_usize() <= w[1].as_usize()), + "get_anti_indices requires ascending input_indices" + ); + + let mut next_unmatched_idx = range.start; + let mut output: Vec<T::Native> = Vec::with_capacity(range.len()); + + for &v in input_indices.values() { + let idx = v.as_usize(); + + if idx < range.start { + continue; + } + if idx >= range.end { + break; + } + + if next_unmatched_idx < idx { + output.extend((next_unmatched_idx..idx).map(|idx| { + T::Native::from_usize(idx).expect("join index exceeds output index type") + })); + } + next_unmatched_idx = idx + 1; + } + + if next_unmatched_idx < range.end { + output.extend((next_unmatched_idx..range.end).map(|idx| { + T::Native::from_usize(idx).expect("join index exceeds output index type") + })); + } + PrimitiveArray::<T>::new(output.into(), None) +} + +/// Returns the intersection of `range` and `input_indices`, omitting duplicates. +/// +/// `input_indices` must be sorted ascending and contain no nulls. +pub(crate) fn get_semi_indices<T: ArrowPrimitiveType>( + range: Range<usize>, + input_indices: &PrimitiveArray<T>, +) -> PrimitiveArray<T> +where + NativeAdapter<T>: From<<T as ArrowPrimitiveType>::Native>, +{ + debug_assert_eq!( + input_indices.null_count(), + 0, + "get_semi_indices requires non-null input_indices" + ); + debug_assert!( + input_indices + .values() + .windows(2) + .all(|w| w[0].as_usize() <= w[1].as_usize()), + "get_semi_indices requires ascending input_indices" + ); + + let mut prev_idx: Option<usize> = None; + let mut output = Vec::with_capacity(input_indices.len().min(range.len())); + + for &v in input_indices.values() { + let idx = v.as_usize(); + + if idx < range.start { + continue; + } + if idx >= range.end { + break; + } + + if prev_idx.replace(idx) != Some(idx) { + output.push(v); + } + } + + PrimitiveArray::<T>::new(output.into(), None) +} + +pub(crate) fn get_mark_indices<T: ArrowPrimitiveType>( + range: &Range<usize>, + input_indices: &PrimitiveArray<T>, +) -> PrimitiveArray<UInt32Type> +where + NativeAdapter<T>: From<<T as ArrowPrimitiveType>::Native>, +{ + let mut bitmap = build_range_bitmap(range, input_indices); + PrimitiveArray::new( + vec![0; range.len()].into(), + Some(NullBuffer::new(bitmap.finish())), + ) +} + +fn build_range_bitmap<T: ArrowPrimitiveType>( + range: &Range<usize>, + input: &PrimitiveArray<T>, +) -> BooleanBufferBuilder { + let mut builder = BooleanBufferBuilder::new(range.len()); + builder.append_n(range.len(), false); + + input.iter().flatten().for_each(|v| { + let idx = v.as_usize(); + if range.contains(&idx) { + builder.set_bit(idx - range.start, true); + } + }); + + builder +} + +/// Appends probe indices in order by considering the given build indices. +/// +/// This function constructs new build and probe indices by iterating through +/// the provided indices, and appends any missing values between previous and +/// current probe index with a corresponding null build index. +/// +/// # Parameters +/// +/// - `build_indices`: `PrimitiveArray` of `UInt64Type` containing build indices. +/// - `probe_indices`: `PrimitiveArray` of `UInt32Type` containing probe indices. +/// - `range`: The range of indices to consider. +/// +/// # Returns +/// +/// A tuple of two arrays: +/// - A `PrimitiveArray` of `UInt64Type` with the newly constructed build indices. +/// - A `PrimitiveArray` of `UInt32Type` with the newly constructed probe indices. +fn append_probe_indices_in_order( + build_indices: &PrimitiveArray<UInt64Type>, + probe_indices: &PrimitiveArray<UInt32Type>, + range: Range<usize>, +) -> (PrimitiveArray<UInt64Type>, PrimitiveArray<UInt32Type>) { + // Builders for new indices: + let mut new_build_indices = UInt64Builder::new(); + let mut new_probe_indices = UInt32Builder::new(); + // Set previous index as the start index for the initial loop: + let mut prev_index = range.start as u32; + // Zip the two iterators. + debug_assert!(build_indices.len() == probe_indices.len()); + for (build_index, probe_index) in build_indices + .values() + .into_iter() + .zip(probe_indices.values()) + { + // Append values between previous and current probe index with null build index: + for value in prev_index..*probe_index { + new_probe_indices.append_value(value); + new_build_indices.append_null(); + } + // Append current indices: + new_probe_indices.append_value(*probe_index); + new_build_indices.append_value(*build_index); + // Set current probe index as previous for the next iteration: + prev_index = probe_index + 1; + } + // Append remaining probe indices after the last valid probe index with null build index. + for value in prev_index..range.end as u32 { + new_probe_indices.append_value(value); + new_build_indices.append_null(); + } + // Build arrays and return: + (new_build_indices.finish(), new_probe_indices.finish()) +} + +/// Metrics for build & probe joins +#[derive(Clone, Debug)] +pub(crate) struct BuildProbeJoinMetrics { + pub(crate) baseline: BaselineMetrics, + /// Total time for collecting build-side of join + pub(crate) build_time: metrics::Time, + /// Number of batches consumed by build-side + pub(crate) build_input_batches: metrics::Count, + /// Number of rows consumed by build-side + pub(crate) build_input_rows: metrics::Count, + /// Memory used by build-side in bytes + pub(crate) build_mem_used: metrics::Gauge, + /// Total time for joining probe-side batches to the build-side batches + pub(crate) join_time: metrics::Time, + /// Number of batches consumed by probe-side of this operator + pub(crate) input_batches: metrics::Count, + /// Number of rows consumed by probe-side this operator + pub(crate) input_rows: metrics::Count, + /// Fraction of probe rows that found more than one match + pub(crate) probe_hit_rate: metrics::RatioMetrics, + /// Average number of build matches per matched probe row + pub(crate) avg_fanout: metrics::RatioMetrics, +} + +// This Drop implementation updates the elapsed compute part of the metrics. +// +// Why is this in a Drop? +// - We keep track of build_time and join_time separately, but baseline metrics have +// a total elapsed_compute time. Instead of remembering to update both the metrics +// at the same time, we chose to update elapsed_compute once at the end - summing up +// both the parts. +// +// How does this work? +// - The elapsed_compute `Time` is represented by an `Arc<AtomicUsize>`. So even when +// this `BuildProbeJoinMetrics` is dropped, the elapsed_compute is usable through the +// Arc reference. +impl Drop for BuildProbeJoinMetrics { + fn drop(&mut self) { + self.baseline.elapsed_compute().add(&self.build_time); + self.baseline.elapsed_compute().add(&self.join_time); + } +} + +impl BuildProbeJoinMetrics { + pub fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + let baseline = BaselineMetrics::new(metrics, partition); + + let join_time = MetricBuilder::new(metrics).subset_time("join_time", partition); + + let build_time = MetricBuilder::new(metrics).subset_time("build_time", partition); + + let build_input_batches = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("build_input_batches", partition); + + let build_input_rows = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("build_input_rows", partition); + + let build_mem_used = + MetricBuilder::new(metrics).peak_memory_usage("build_mem_used", partition); + + let input_batches = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("input_batches", partition); + + let input_rows = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("input_rows", partition); + + let probe_hit_rate = MetricBuilder::new(metrics) + .with_type(MetricType::Summary) + .ratio_metrics("probe_hit_rate", partition); + + let avg_fanout = MetricBuilder::new(metrics) + .with_type(MetricType::Summary) + .ratio_metrics("avg_fanout", partition); + + Self { + build_time, + build_input_batches, + build_input_rows, + build_mem_used, + join_time, + input_batches, + input_rows, + baseline, + probe_hit_rate, + avg_fanout, + } + } +} + +/// The `handle_state` macro is designed to process the result of a state-changing +/// operation. It operates on a `StatefulStreamResult` by matching its variants and +/// executing corresponding actions. This macro is used to streamline code that deals +/// with state transitions, reducing boilerplate and improving readability. +/// +/// # Cases +/// +/// - `Ok(StatefulStreamResult::Continue)`: Continues the loop, indicating the +/// stream join operation should proceed to the next step. +/// - `Ok(StatefulStreamResult::Ready(result))`: Returns a `Poll::Ready` with the +/// result, either yielding a value or indicating the stream is awaiting more +/// data. +/// - `Err(e)`: Returns a `Poll::Ready` containing an error, signaling an issue +/// during the stream join operation. +/// +/// # Arguments +/// +/// * `$match_case`: An expression that evaluates to a `Result<StatefulStreamResult<_>>`. +#[macro_export] +macro_rules! handle_state { + ($match_case:expr) => { + match $match_case { + Ok(StatefulStreamResult::Continue) => continue, + Ok(StatefulStreamResult::Ready(result)) => { + Poll::Ready(Ok(result).transpose()) + } + Err(e) => Poll::Ready(Some(Err(e))), + } + }; +} + +/// Represents the result of a stateful operation. +/// +/// This enumeration indicates whether the state produced a result that is +/// ready for use (`Ready`) or if the operation requires continuation (`Continue`). +/// +/// Variants: +/// - `Ready(T)`: Indicates that the operation is complete with a result of type `T`. +/// - `Continue`: Indicates that the operation is not yet complete and requires further +/// processing or more data. When this variant is returned, it typically means that the +/// current invocation of the state did not produce a final result, and the operation +/// should be invoked again later with more data and possibly with a different state. +pub enum StatefulStreamResult<T> { + Ready(T), + Continue, +} + +pub(crate) fn symmetric_join_output_partitioning( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + join_type: &JoinType, +) -> Result<Partitioning> { + let left_columns_len = left.schema().fields.len(); + let left_partitioning = left.output_partitioning(); + let right_partitioning = right.output_partitioning(); + let result = match join_type { + JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + left_partitioning.clone() + } + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + right_partitioning.clone() + } + JoinType::Inner | JoinType::Right => { + adjust_right_output_partitioning(right_partitioning, left_columns_len)? + } + JoinType::Full => { + // We could also use left partition count as they are necessarily equal. + Partitioning::UnknownPartitioning(right_partitioning.partition_count()) + } + }; + Ok(result) +} + +pub(crate) fn asymmetric_join_output_partitioning( + left: &Arc<dyn ExecutionPlan>, + right: &Arc<dyn ExecutionPlan>, + join_type: &JoinType, +) -> Result<Partitioning> { + let result = match join_type { + JoinType::Inner | JoinType::Right => adjust_right_output_partitioning( + right.output_partitioning(), + left.schema().fields().len(), + )?, + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + right.output_partitioning().clone() + } + JoinType::Left + | JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::Full + | JoinType::LeftMark => Partitioning::UnknownPartitioning( + right.output_partitioning().partition_count(), + ), + }; + Ok(result) +} + +/// Trait for incrementally generating Join output. +/// +/// This trait is used to limit some join outputs +/// so it does not produce single large batches +pub(crate) trait BatchTransformer: Debug + Clone { + /// Sets the next `RecordBatch` to be processed. + fn set_batch(&mut self, batch: RecordBatch); + + /// Retrieves the next `RecordBatch` from the transformer. + /// Returns `None` if all batches have been produced. + /// The boolean flag indicates whether the batch is the last one. + fn next(&mut self) -> Option<(RecordBatch, bool)>; +} + +#[derive(Debug, Clone)] +/// A batch transformer that does nothing. +pub(crate) struct NoopBatchTransformer { + /// RecordBatch to be processed + batch: Option<RecordBatch>, +} + +impl NoopBatchTransformer { + pub fn new() -> Self { + Self { batch: None } + } +} + +impl BatchTransformer for NoopBatchTransformer { + fn set_batch(&mut self, batch: RecordBatch) { + self.batch = Some(batch); + } + + fn next(&mut self) -> Option<(RecordBatch, bool)> { + self.batch.take().map(|batch| (batch, true)) + } +} + +#[derive(Debug, Clone)] +/// Splits large batches into smaller batches with a maximum number of rows. +pub(crate) struct BatchSplitter { + /// RecordBatch to be split + batch: Option<RecordBatch>, + /// Maximum number of rows in a split batch + batch_size: usize, + /// Current row index + row_index: usize, +} + +impl BatchSplitter { + /// Creates a new `BatchSplitter` with the specified batch size. + pub(crate) fn new(batch_size: usize) -> Self { + Self { + batch: None, + batch_size, + row_index: 0, + } + } +} + +impl BatchTransformer for BatchSplitter { + fn set_batch(&mut self, batch: RecordBatch) { + self.batch = Some(batch); + self.row_index = 0; + } + + fn next(&mut self) -> Option<(RecordBatch, bool)> { + let Some(batch) = &self.batch else { + return None; + }; + + let remaining_rows = batch.num_rows() - self.row_index; + let rows_to_slice = remaining_rows.min(self.batch_size); + let sliced_batch = batch.slice(self.row_index, rows_to_slice); + self.row_index += rows_to_slice; + + let mut last = false; + if self.row_index >= batch.num_rows() { + self.batch = None; + last = true; + } + + Some((sliced_batch, last)) + } +} + +/// When the order of the join inputs are changed, the output order of columns +/// must remain the same. +/// +/// Joins output columns from their left input followed by their right input. +/// Thus if the inputs are reordered, the output columns must be reordered to +/// match the original order. +pub fn reorder_output_after_swap( + plan: Arc<dyn ExecutionPlan>, + left_schema: &Schema, + right_schema: &Schema, +) -> Result<Arc<dyn ExecutionPlan>> { + let proj = ProjectionExec::try_new( + swap_reverting_projection(left_schema, right_schema), + plan, + )?; + Ok(Arc::new(proj)) +} + +/// When the order of the join is changed, the output order of columns must +/// remain the same. +/// +/// Returns the expressions that will allow to swap back the values from the +/// original left as the first columns and those on the right next. +fn swap_reverting_projection( + left_schema: &Schema, + right_schema: &Schema, +) -> Vec<ProjectionExpr> { + let right_cols = + right_schema + .fields() + .iter() + .enumerate() + .map(|(i, f)| ProjectionExpr { + expr: Arc::new(Column::new(f.name(), i)) as Arc<dyn PhysicalExpr>, + alias: f.name().to_owned(), + }); + let right_len = right_cols.len(); + let left_cols = + left_schema + .fields() + .iter() + .enumerate() + .map(|(i, f)| ProjectionExpr { + expr: Arc::new(Column::new(f.name(), right_len + i)) + as Arc<dyn PhysicalExpr>, + alias: f.name().to_owned(), + }); + + left_cols.chain(right_cols).collect() +} + +/// This function swaps the given join's projection. +pub fn swap_join_projection( + left_schema_len: usize, + right_schema_len: usize, + projection: Option<&[usize]>, + join_type: &JoinType, +) -> Option<Vec<usize>> { + match join_type { + // For Anti/Semi join types, projection should remain unmodified, + // since these joins output schema remains the same after swap + JoinType::LeftAnti + | JoinType::LeftSemi + | JoinType::RightAnti + | JoinType::RightSemi + | JoinType::LeftMark + | JoinType::RightMark => projection.map(|p| p.to_vec()), + _ => projection.map(|p| { + p.iter() + .map(|i| { + // If the index is less than the left schema length, it is from + // the left schema, so we add the right schema length to it. + // Otherwise, it is from the right schema, so we subtract the left + // schema length from it. + if *i < left_schema_len { + *i + right_schema_len + } else { + *i - left_schema_len + } + }) + .collect() + }), + } +} + +/// Updates `hash_map` with new entries from `batch` evaluated against the expressions `on` +/// using `offset` as a start value for `batch` row indices. +/// +/// `fifo_hashmap` sets the order of iteration over `batch` rows while updating hashmap, +/// which allows to keep either first (if set to true) or last (if set to false) row index +/// as a chain head for rows with equal hash values. +/// +/// Under [`NullEquality::NullEqualsNothing`], rows with a NULL in any key +/// column can never match a probe row, so they are not inserted into the map. +#[expect(clippy::too_many_arguments)] +pub fn update_hash( + on: &[PhysicalExprRef], + batch: &RecordBatch, + hash_map: &mut dyn JoinHashMapType, + offset: usize, + random_state: &RandomState, + hashes_buffer: &mut [u64], + deleted_offset: usize, + fifo_hashmap: bool, + null_equality: NullEquality, +) -> Result<()> { + // evaluate the keys + let keys_values = evaluate_expressions_to_arrays(on, batch)?; + + // calculate the hash values + let hash_values = create_hashes(&keys_values, random_state, hashes_buffer)?; + + // For usual JoinHashmap, the implementation is void. + hash_map.extend_zero(batch.num_rows()); + + // Unmatchable NULL-key rows are filtered out below. + let valid_keys = matchable_join_keys(&keys_values, null_equality); + + // Updating JoinHashMap from hash values iterator + let hash_values_iter = hash_values + .iter() + .enumerate() + .filter(|(i, _)| valid_keys.as_ref().is_none_or(|nulls| nulls.is_valid(*i))) + .map(|(i, val)| (i + offset, val)); + + if fifo_hashmap { + hash_map.update_from_iter(Box::new(hash_values_iter.rev()), deleted_offset); + } else { + hash_map.update_from_iter(Box::new(hash_values_iter), deleted_offset); + } + + Ok(()) +} + +/// Returns the combined validity of the join key columns `join_key_arrays`: a row +/// is valid only if every key column is non-NULL at that row. +/// +/// Returns `None` when no rows need to be filtered: either every row has +/// fully non-NULL keys, or `null_equality` is +/// [`NullEquality::NullEqualsNull`], where NULL keys are matchable. +pub(crate) fn matchable_join_keys( + join_key_arrays: &[ArrayRef], + null_equality: NullEquality, +) -> Option<NullBuffer> { + match null_equality { + NullEquality::NullEqualsNothing => { + let logical_nulls: Vec<_> = join_key_arrays + .iter() + .map(|values| values.logical_nulls()) + .collect(); + NullBuffer::union_many(logical_nulls.iter().map(Option::as_ref)) + // An all-valid array can still have a validity buffer; return + // `None` in that case, since there is nothing to filter. + .filter(|nulls| nulls.null_count() > 0) + } + NullEquality::NullEqualsNull => None, + } +} + +pub(super) fn equal_rows_arr( + indices_left: &UInt64Array, + indices_right: &UInt32Array, + left_arrays: &[ArrayRef], + right_arrays: &[ArrayRef], + null_equality: NullEquality, +) -> Result<(UInt64Array, UInt32Array)> { + if indices_left.len() != indices_right.len() { + return Err(internal_datafusion_err!( + "Cannot compare join indices with different lengths: left={}, right={}", + indices_left.len(), + indices_right.len() + )); + } + + if left_arrays.len() != right_arrays.len() { + return Err(internal_datafusion_err!( + "Cannot compare join keys with different column counts: left={}, right={}", + left_arrays.len(), + right_arrays.len() + )); + } + + if left_arrays.is_empty() { + return Ok((Vec::<u64>::new().into(), Vec::<u32>::new().into())); + } + + // Fast path: single-column keys of a specialized type run a monomorphized + // equality loop, avoiding the per-pair boxed `DynComparator` dispatch and + // `Ordering` computation of the general `JoinKeyComparator` path. Falls + // through to the general path for multi-column keys and unspecialized + // types (e.g. floats, dictionaries, nested). + let single_col_fast_path = if left_arrays.len() == 1 { + equal_rows_single_col( + indices_left, + indices_right, + left_arrays[0].as_ref(), + right_arrays[0].as_ref(), + null_equality, + ) + } else { + None + }; + if let Some(res) = single_col_fast_path { + return Ok(res); + } + + let sort_options = vec![SortOptions::default(); left_arrays.len()]; + let comparator = + JoinKeyComparator::new(left_arrays, right_arrays, &sort_options, null_equality)?; + + let mut left_filtered = Vec::with_capacity(indices_left.len()); + let mut right_filtered = Vec::with_capacity(indices_right.len()); + + for (left, right) in indices_left.values().iter().zip(indices_right.values()) { + let left_idx = usize::try_from(*left).map_err(|_| { + internal_datafusion_err!("Join index {left} can not be represented as usize") + })?; + let right_idx = *right as usize; + + if comparator.is_equal(left_idx, right_idx) { + left_filtered.push(*left); + right_filtered.push(*right); + } + } + + Ok((left_filtered.into(), right_filtered.into())) +} + +/// Specialized single-column equi-join key filtering. +/// +/// Dispatches once on the key column's type and runs a monomorphized equality +/// loop with typed value comparison. This avoids the per-pair boxed +/// `DynComparator` call and the three-way `Ordering` computation used by the +/// general [`JoinKeyComparator`] path, which dominates for high-fanout +/// single-column joins (e.g. long string keys with near-100% match rates). +/// +/// Returns `None` for types it does not specialize (including when the left and +/// right key types differ, handled by the failed downcast) so the caller falls +/// back to the general path. Floats are intentionally excluded so their `-0.0` / +/// `NaN` semantics stay on the exact same code path as before. +fn equal_rows_single_col( + indices_left: &UInt64Array, + indices_right: &UInt32Array, + left: &dyn Array, + right: &dyn Array, + null_equality: NullEquality, +) -> Option<(UInt64Array, UInt32Array)> { + let null_equals_null = matches!(null_equality, NullEquality::NullEqualsNull); + + macro_rules! eq_loop { + ($T:ty) => {{ + let l = left.as_any().downcast_ref::<$T>()?; + let r = right.as_any().downcast_ref::<$T>()?; + + let mut left_filtered = Vec::with_capacity(indices_left.len()); + let mut right_filtered = Vec::with_capacity(indices_right.len()); + + for (left_idx, right_idx) in + indices_left.values().iter().zip(indices_right.values()) + { + let i = *left_idx as usize; + let j = *right_idx as usize; + + let is_equal = match (l.is_null(i), r.is_null(j)) { + (false, false) => l.value(i) == r.value(j), + (true, true) => null_equals_null, + _ => false, + }; + + if is_equal { + left_filtered.push(*left_idx); + right_filtered.push(*right_idx); + } + } + + return Some((left_filtered.into(), right_filtered.into())); + }}; + } + + match left.data_type() { + DataType::Boolean => eq_loop!(BooleanArray), + DataType::Int8 => eq_loop!(Int8Array), + DataType::Int16 => eq_loop!(Int16Array), + DataType::Int32 => eq_loop!(Int32Array), + DataType::Int64 => eq_loop!(Int64Array), + DataType::UInt8 => eq_loop!(UInt8Array), + DataType::UInt16 => eq_loop!(UInt16Array), + DataType::UInt32 => eq_loop!(UInt32Array), + DataType::UInt64 => eq_loop!(UInt64Array), + DataType::Decimal128(..) => eq_loop!(Decimal128Array), + DataType::Binary => eq_loop!(BinaryArray), + DataType::LargeBinary => eq_loop!(LargeBinaryArray), + DataType::BinaryView => eq_loop!(BinaryViewArray), + DataType::FixedSizeBinary(_) => eq_loop!(FixedSizeBinaryArray), + DataType::Utf8 => eq_loop!(StringArray), + DataType::LargeUtf8 => eq_loop!(LargeStringArray), + DataType::Utf8View => eq_loop!(StringViewArray), + DataType::Date32 => eq_loop!(Date32Array), + DataType::Date64 => eq_loop!(Date64Array), + DataType::Timestamp(time_unit, _) => match time_unit { + TimeUnit::Second => eq_loop!(TimestampSecondArray), + TimeUnit::Millisecond => eq_loop!(TimestampMillisecondArray), + TimeUnit::Microsecond => eq_loop!(TimestampMicrosecondArray), + TimeUnit::Nanosecond => eq_loop!(TimestampNanosecondArray), + }, + _ => None, + } +} + +/// Pre-built comparator for join key columns that eliminates per-row type +/// dispatch. Wraps `arrow_ord::ord::DynComparator` closures built once per +/// batch pair, used for all row comparisons within those batches. +/// +/// The first key column is stored separately so that single-column joins +/// (the common case) avoid Vec iteration entirely, and multi-column joins +/// short-circuit without entering the loop when the first column is +/// selective. +/// +/// Null handling is baked into the closures at construction time: +/// - `NullEqualsNull`: `make_comparator` returns `Equal` for both-null, which +/// is the desired behavior. Closures are used as-is. +/// - `NullEqualsNothing`: columns where both sides contain nulls get a wrapper +/// that returns `Less` for both-null. Columns where one side has no nulls +/// skip the wrapper since both-null is impossible. +/// +/// Because `NullEqualsNothing` wraps comparators to return `Less` for +/// both-null, `is_equal` will return `false` for both-null rows when that +/// mode is active. Callers needing both-null == equal semantics (e.g., +/// buffered head/tail equality in SMJ) should construct with +/// `NullEqualsNull`. +pub struct JoinKeyComparator { + first: DynComparator, + rest: Vec<DynComparator>, +} + +impl JoinKeyComparator { + /// Build comparators for each join key column pair. + pub fn new( + left_arrays: &[ArrayRef], + right_arrays: &[ArrayRef], + sort_options: &[SortOptions], + null_equality: NullEquality, + ) -> Result<Self> { + debug_assert_eq!(left_arrays.len(), right_arrays.len()); + debug_assert_eq!(left_arrays.len(), sort_options.len()); + + let mut iter = left_arrays + .iter() + .zip(right_arrays.iter()) + .zip(sort_options.iter()) + .map(|((l, r), opts)| { + // `make_comparator` uses IEEE 754 totalOrder for floats and + // treats `-0.0` / `+0.0` as distinct. Normalize float arrays + // so SMJ / piecewise-merge equi-keys honor SQL equality; + // no-op (Arc::clone) for non-floats and for float arrays + // that contain no `-0.0`. `normalize_float_zero` preserves + // null positions, so the original null masks below remain + // valid. + let l_norm = normalize_float_zero(l); + let r_norm = normalize_float_zero(r); + let inner = make_comparator(l_norm.as_ref(), r_norm.as_ref(), *opts)?; + if null_equality == NullEquality::NullEqualsNothing { + let ln = l.logical_nulls().filter(|n| n.null_count() > 0); + let rn = r.logical_nulls().filter(|n| n.null_count() > 0); + match (ln, rn) { + // Both sides have nulls — wrap to override both-null. + (Some(ln), Some(rn)) => Ok(Box::new(move |i, j| { + if ln.is_null(i) && rn.is_null(j) { + Ordering::Less + } else { + inner(i, j) + } + }) + as DynComparator), + // One side has no nulls — both-null impossible, no wrap. + _ => Ok(inner), + } + } else { + Ok(inner) + } + }); + + let first = iter.next().expect("join must have at least one key")?; + let rest = iter.collect::<Result<Vec<_>>>()?; + Ok(Self { first, rest }) + } + + /// Compare row `left` (in the left arrays) with row `right` (in the right + /// arrays). Returns the lexicographic ordering across all key columns. + #[inline] + pub fn compare(&self, left: usize, right: usize) -> Ordering { + let ord = (self.first)(left, right); + if ord != Ordering::Equal || self.rest.is_empty() { + return ord; + } + for cmp_fn in &self.rest { + let ord = cmp_fn(left, right); + if ord != Ordering::Equal { + return ord; + } + } + Ordering::Equal + } + + /// Check equality of row `left` (in the left arrays) with row `right` + /// (in the right arrays). Both-null is treated as equal when constructed + /// with `NullEqualsNull`. With `NullEqualsNothing`, both-null returns + /// `false` because the override is baked into the comparators. + #[inline] + pub fn is_equal(&self, left: usize, right: usize) -> bool { + if (self.first)(left, right) != Ordering::Equal { + return false; + } + for cmp_fn in &self.rest { + if cmp_fn(left, right) != Ordering::Equal { + return false; + } + } + true + } +} + +/// Get comparison result of two rows of join arrays +pub fn compare_join_arrays( + left_arrays: &[ArrayRef], + left: usize, + right_arrays: &[ArrayRef], + right: usize, + sort_options: &[SortOptions], + null_equality: NullEquality, +) -> Result<Ordering> { + let mut res = Ordering::Equal; + for ((left_array, right_array), sort_options) in + left_arrays.iter().zip(right_arrays).zip(sort_options) + { + macro_rules! compare_value { + ($T:ty) => {{ + let left_array = left_array.as_any().downcast_ref::<$T>().unwrap(); + let right_array = right_array.as_any().downcast_ref::<$T>().unwrap(); + match (left_array.is_null(left), right_array.is_null(right)) { + (false, false) => { + let left_value = &left_array.value(left); + let right_value = &right_array.value(right); + res = left_value.partial_cmp(right_value).unwrap(); + if sort_options.descending { + res = res.reverse(); + } + } + (true, false) => { + res = if sort_options.nulls_first { + Ordering::Less + } else { + Ordering::Greater + }; + } + (false, true) => { + res = if sort_options.nulls_first { + Ordering::Greater + } else { + Ordering::Less + }; + } + _ => { + res = match null_equality { + NullEquality::NullEqualsNothing => Ordering::Less, + NullEquality::NullEqualsNull => Ordering::Equal, + }; + } + } + }}; + } + + match left_array.data_type() { + DataType::Null => {} + DataType::Boolean => compare_value!(BooleanArray), + DataType::Int8 => compare_value!(Int8Array), + DataType::Int16 => compare_value!(Int16Array), + DataType::Int32 => compare_value!(Int32Array), + DataType::Int64 => compare_value!(Int64Array), + DataType::UInt8 => compare_value!(UInt8Array), + DataType::UInt16 => compare_value!(UInt16Array), + DataType::UInt32 => compare_value!(UInt32Array), + DataType::UInt64 => compare_value!(UInt64Array), + DataType::Float32 => compare_value!(Float32Array), + DataType::Float64 => compare_value!(Float64Array), + DataType::Binary => compare_value!(BinaryArray), + DataType::BinaryView => compare_value!(BinaryViewArray), + DataType::FixedSizeBinary(_) => compare_value!(FixedSizeBinaryArray), + DataType::LargeBinary => compare_value!(LargeBinaryArray), + DataType::Utf8 => compare_value!(StringArray), + DataType::Utf8View => compare_value!(StringViewArray), + DataType::LargeUtf8 => compare_value!(LargeStringArray), + DataType::Decimal128(..) => compare_value!(Decimal128Array), + DataType::Timestamp(time_unit, None) => match time_unit { + TimeUnit::Second => compare_value!(TimestampSecondArray), + TimeUnit::Millisecond => compare_value!(TimestampMillisecondArray), + TimeUnit::Microsecond => compare_value!(TimestampMicrosecondArray), + TimeUnit::Nanosecond => compare_value!(TimestampNanosecondArray), + }, + DataType::Date32 => compare_value!(Date32Array), + DataType::Date64 => compare_value!(Date64Array), + dt => { + return not_impl_err!( + "Unsupported data type in sort merge join comparator: {}", + dt + ); + } + } + if !res.is_eq() { + break; + } + } + Ok(res) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::pin::Pin; + + use super::*; + + use arrow::datatypes::{DataType, Fields}; + use arrow::error::{ArrowError, Result as ArrowResult}; + use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; + use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err}; + use datafusion_physical_expr::PhysicalSortExpr; + + use rstest::rstest; + + fn assert_u32_values(array: &UInt32Array, expected: &[u32]) { + assert_eq!(array.values().as_ref(), expected); + } + + #[test] + fn get_anti_indices_returns_unmatched_range_indices() { + let input = UInt32Array::from(vec![3, 5, 5]); + + let result = get_anti_indices(2..8, &input); + + assert_u32_values(&result, &[2, 4, 6, 7]); + } + + #[test] + fn get_anti_indices_ignores_out_of_range_indices() { + let input = UInt32Array::from(vec![0, 1, 3, 5, 8, 12]); + + let result = get_anti_indices(2..8, &input); + + assert_u32_values(&result, &[2, 4, 6, 7]); + } + + #[test] + fn update_hash_skips_null_keys_for_null_equals_nothing() -> Result<()> { + use crate::joins::join_hash_map::JoinHashMapU32; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(1), + ]))], + )?; + let on: Vec<PhysicalExprRef> = vec![Arc::new(Column::new("a", 0))]; + let random_state = RandomState::with_seed(42); + let mut hashes_buffer = vec![0; batch.num_rows()]; + create_hashes([batch.column(0)], &random_state, &mut hashes_buffer)?; + + let matched_build_indices = + |map: &JoinHashMapU32, hashes_buffer: &[u64]| -> Vec<u64> { + let mut input_indices = vec![]; + let mut match_indices = vec![]; + map.get_matched_indices_with_limit_offset( + hashes_buffer, + None, + 8192, + (0, None), + &mut input_indices, + &mut match_indices, + ); + match_indices.sort_unstable(); + match_indices.dedup(); + match_indices + }; + + let mut map = JoinHashMapU32::with_capacity(batch.num_rows()); + update_hash( + &on, + &batch, + &mut map, + 0, + &random_state, + &mut hashes_buffer, + 0, + true, + NullEquality::NullEqualsNothing, + )?; + // NULL keys can never match under NullEqualsNothing, so they must not + // be inserted into the map. Assert row indices rather than map length: + // with forced hash collisions, multiple logical keys can share one + // hash table entry. + assert_eq!(matched_build_indices(&map, &hashes_buffer), vec![0, 2, 4]); + + let mut map = JoinHashMapU32::with_capacity(batch.num_rows()); + update_hash( + &on, + &batch, + &mut map, + 0, + &random_state, + &mut hashes_buffer, + 0, + true, + NullEquality::NullEqualsNull, + )?; + // Under NullEqualsNull, NULL keys can match, so the build-side NULL + // rows must be present in the map. + assert_eq!( + matched_build_indices(&map, &hashes_buffer), + vec![0, 1, 2, 3, 4] + ); + + Ok(()) + } + + #[test] + fn get_anti_indices_handles_dense_matches() { + let input = UInt32Array::from(vec![2, 3, 4, 5]); + + let result = get_anti_indices(2..6, &input); + + assert!(result.is_empty()); + } + + #[test] + fn get_anti_indices_handles_sparse_matches() { + let input = UInt32Array::from(vec![0, 8]); + + let result = get_anti_indices(2..6, &input); + + assert_u32_values(&result, &[2, 3, 4, 5]); + } + + #[test] + fn get_semi_indices_returns_distinct_matches_in_range() { + let input = UInt32Array::from(vec![1, 3, 3, 3, 5, 8]); + + let result = get_semi_indices(2..7, &input); + + assert_u32_values(&result, &[3, 5]); + } + + #[test] + fn get_semi_indices_ignores_out_of_range_indices() { + let input = UInt32Array::from(vec![0, 1, 3, 5, 8, 12]); + + let result = get_semi_indices(2..8, &input); + + assert_u32_values(&result, &[3, 5]); + } + + #[test] + fn get_semi_indices_handles_dense_matches() { + let input = UInt32Array::from(vec![2, 3, 4, 5]); + + let result = get_semi_indices(2..6, &input); + + assert_u32_values(&result, &[2, 3, 4, 5]); + } + + #[test] + fn get_semi_indices_handles_empty_input() { + let input = UInt32Array::from(Vec::<u32>::new()); + + let result = get_semi_indices(2..6, &input); + + assert!(result.is_empty()); + } + + fn check( + left: &[Column], + right: &[Column], + on: &[(PhysicalExprRef, PhysicalExprRef)], + ) -> Result<()> { + let left = left + .iter() + .map(|x| x.to_owned()) + .collect::<HashSet<Column>>(); + let right = right + .iter() + .map(|x| x.to_owned()) + .collect::<HashSet<Column>>(); + check_join_set_is_valid(&left, &right, on) + } + + #[test] + fn check_valid() -> Result<()> { + let left = vec![Column::new("a", 0), Column::new("b1", 1)]; + let right = vec![Column::new("a", 0), Column::new("b2", 1)]; + let on = &[( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )]; + + check(&left, &right, on)?; + Ok(()) + } + + #[test] + fn check_not_in_right() { + let left = vec![Column::new("a", 0), Column::new("b", 1)]; + let right = vec![Column::new("b", 0)]; + let on = &[( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )]; + + assert!(check(&left, &right, on).is_err()); + } + + #[tokio::test] + async fn check_error_nesting() { + let once_fut = OnceFut::<()>::new(async { + arrow_err!(ArrowError::CsvError("some error".to_string())) + }); + + struct TestFut(OnceFut<()>); + impl Future for TestFut { + type Output = ArrowResult<()>; + + fn poll( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Self::Output> { + match ready!(self.0.get(cx)) { + Ok(()) => Poll::Ready(Ok(())), + Err(e) => Poll::Ready(Err(e.into())), + } + } + } + + let res = TestFut(once_fut).await; + let arrow_err_from_fut = res.expect_err("once_fut always return error"); + + let wrapped_err = DataFusionError::from(arrow_err_from_fut); + let root_err = wrapped_err.find_root(); + + let _expected = + arrow_datafusion_err!(ArrowError::CsvError("some error".to_owned())); + + assert!(matches!(root_err, _expected)) + } + + #[test] + fn check_not_in_left() { + let left = vec![Column::new("b", 0)]; + let right = vec![Column::new("a", 0)]; + let on = &[( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )]; + + assert!(check(&left, &right, on).is_err()); + } + + #[test] + fn check_collision() { + // column "a" would appear both in left and right + let left = vec![Column::new("a", 0), Column::new("c", 1)]; + let right = vec![Column::new("a", 0), Column::new("b", 1)]; + let on = &[( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("b", 1)) as _, + )]; + + assert!(check(&left, &right, on).is_ok()); + } + + #[test] + fn check_in_right() { + let left = vec![Column::new("a", 0), Column::new("c", 1)]; + let right = vec![Column::new("b", 0)]; + let on = &[( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("b", 0)) as _, + )]; + + assert!(check(&left, &right, on).is_ok()); + } + + #[test] + fn test_join_schema() -> Result<()> { + let a = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let a_nulls = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let b = Schema::new(vec![Field::new("b", DataType::Int32, false)]); + let b_nulls = Schema::new(vec![Field::new("b", DataType::Int32, true)]); + + let cases = vec![ + (&a, &b, JoinType::Inner, &a, &b), + (&a, &b_nulls, JoinType::Inner, &a, &b_nulls), + (&a_nulls, &b, JoinType::Inner, &a_nulls, &b), + (&a_nulls, &b_nulls, JoinType::Inner, &a_nulls, &b_nulls), + // right input of a `LEFT` join can be null, regardless of input nullness + (&a, &b, JoinType::Left, &a, &b_nulls), + (&a, &b_nulls, JoinType::Left, &a, &b_nulls), + (&a_nulls, &b, JoinType::Left, &a_nulls, &b_nulls), + (&a_nulls, &b_nulls, JoinType::Left, &a_nulls, &b_nulls), + // left input of a `RIGHT` join can be null, regardless of input nullness + (&a, &b, JoinType::Right, &a_nulls, &b), + (&a, &b_nulls, JoinType::Right, &a_nulls, &b_nulls), + (&a_nulls, &b, JoinType::Right, &a_nulls, &b), + (&a_nulls, &b_nulls, JoinType::Right, &a_nulls, &b_nulls), + // Either input of a `FULL` join can be null + (&a, &b, JoinType::Full, &a_nulls, &b_nulls), + (&a, &b_nulls, JoinType::Full, &a_nulls, &b_nulls), + (&a_nulls, &b, JoinType::Full, &a_nulls, &b_nulls), + (&a_nulls, &b_nulls, JoinType::Full, &a_nulls, &b_nulls), + ]; + + for (left_in, right_in, join_type, left_out, right_out) in cases { + let (schema, _) = build_join_schema(left_in, right_in, &join_type); + + let expected_fields = left_out + .fields() + .iter() + .cloned() + .chain(right_out.fields().iter().cloned()) + .collect::<Fields>(); + + let expected_schema = Schema::new(expected_fields); + assert_eq!( + schema, + expected_schema, + "Mismatch with left_in={}:{}, right_in={}:{}, join_type={:?}", + left_in.fields()[0].name(), + left_in.fields()[0].is_nullable(), + right_in.fields()[0].name(), + right_in.fields()[0].is_nullable(), + join_type + ); + } + + Ok(()) + } + + fn create_stats( + num_rows: Option<usize>, + column_stats: Vec<ColumnStatistics>, + is_exact: bool, + ) -> Statistics { + Statistics { + num_rows: if is_exact { + num_rows.map(Exact) + } else { + num_rows.map(Inexact) + } + .unwrap_or(Absent), + column_statistics: column_stats, + total_byte_size: Absent, + } + } + + fn create_column_stats( + min: Precision<i64>, + max: Precision<i64>, + distinct_count: Precision<usize>, + null_count: Precision<usize>, + ) -> ColumnStatistics { + ColumnStatistics { + distinct_count, + min_value: min.map(ScalarValue::from), + max_value: max.map(ScalarValue::from), + sum_value: Absent, + null_count, + byte_size: Absent, + } + } + + type PartialStats = ( + usize, + Precision<i64>, + Precision<i64>, + Precision<usize>, + Precision<usize>, + ); + + // This is mainly for validating the all edge cases of the estimation, but + // more advanced (and real world test cases) are below where we need some control + // over the expected output (since it depends on join type to join type). + #[test] + fn test_inner_join_cardinality_single_column() -> Result<()> { + let cases: Vec<(PartialStats, PartialStats, Option<Precision<usize>>)> = vec![ + // ------------------------------------------------ + // | left(rows, min, max, distinct, null_count), | + // | right(rows, min, max, distinct, null_count), | + // | expected, | + // ------------------------------------------------ + + // Cardinality computation + // ======================= + // + // distinct(left) == NaN, distinct(right) == NaN + ( + (10, Inexact(1), Inexact(10), Absent, Absent), + (10, Inexact(1), Inexact(10), Absent, Absent), + Some(Inexact(10)), + ), + // range(left) > range(right) + ( + (10, Inexact(6), Inexact(10), Absent, Absent), + (10, Inexact(8), Inexact(10), Absent, Absent), + Some(Inexact(20)), + ), + // range(right) > range(left) + ( + (10, Inexact(8), Inexact(10), Absent, Absent), + (10, Inexact(6), Inexact(10), Absent, Absent), + Some(Inexact(20)), + ), + // range(left) > len(left), range(right) > len(right) + ( + (10, Inexact(1), Inexact(15), Absent, Absent), + (20, Inexact(1), Inexact(40), Absent, Absent), + Some(Inexact(10)), + ), + // Distinct count matches the range + ( + (10, Inexact(1), Inexact(10), Inexact(10), Absent), + (10, Inexact(1), Inexact(10), Inexact(10), Absent), + Some(Inexact(10)), + ), + // Distinct count takes precedence over the range + ( + (10, Inexact(1), Inexact(3), Inexact(10), Absent), + (10, Inexact(1), Inexact(3), Inexact(10), Absent), + Some(Inexact(10)), + ), + // distinct(left) > distinct(right) + ( + (10, Inexact(1), Inexact(10), Inexact(5), Absent), + (10, Inexact(1), Inexact(10), Inexact(2), Absent), + Some(Inexact(20)), + ), + // distinct(right) > distinct(left) + ( + (10, Inexact(1), Inexact(10), Inexact(2), Absent), + (10, Inexact(1), Inexact(10), Inexact(5), Absent), + Some(Inexact(20)), + ), + // min(left) < 0 (range(left) > range(right)) + ( + (10, Inexact(-5), Inexact(5), Absent, Absent), + (10, Inexact(1), Inexact(5), Absent, Absent), + Some(Inexact(10)), + ), + // min(right) < 0, max(right) < 0 (range(right) > range(left)) + ( + (10, Inexact(-25), Inexact(-20), Absent, Absent), + (10, Inexact(-25), Inexact(-15), Absent, Absent), + Some(Inexact(10)), + ), + // range(left) < 0, range(right) >= 0 + // (there isn't a case where both left and right ranges are negative + // so one of them is always going to work, this just proves negative + // ranges with bigger absolute values are not are not accidentally used). + ( + (10, Inexact(-10), Inexact(0), Absent, Absent), + (10, Inexact(0), Inexact(10), Inexact(5), Absent), + Some(Inexact(10)), + ), + // range(left) = 1, range(right) = 1 + ( + (10, Inexact(1), Inexact(1), Absent, Absent), + (10, Inexact(1), Inexact(1), Absent, Absent), + Some(Inexact(100)), + ), + // + // Edge cases + // ========== + // + // No column level stats, fall back to row count. + ( + (10, Absent, Absent, Absent, Absent), + (10, Absent, Absent, Absent, Absent), + Some(Inexact(10)), + ), + // No min or max (or both), but distinct available. + ( + (10, Absent, Absent, Inexact(3), Absent), + (10, Absent, Absent, Inexact(3), Absent), + Some(Inexact(33)), + ), + ( + (10, Inexact(2), Absent, Inexact(3), Absent), + (10, Absent, Inexact(5), Inexact(3), Absent), + Some(Inexact(33)), + ), + ( + (10, Absent, Inexact(3), Inexact(3), Absent), + (10, Inexact(1), Absent, Inexact(3), Absent), + Some(Inexact(33)), + ), + // No min or max, fall back to row count + ( + (10, Absent, Inexact(3), Absent, Absent), + (10, Inexact(1), Absent, Absent, Absent), + Some(Inexact(10)), + ), + // Non overlapping min/max (when exact=False). + ( + (10, Absent, Inexact(4), Absent, Absent), + (10, Inexact(5), Absent, Absent, Absent), + Some(Inexact(0)), + ), + ( + (10, Inexact(0), Inexact(10), Absent, Absent), + (10, Inexact(11), Inexact(20), Absent, Absent), + Some(Inexact(0)), + ), + ( + (10, Inexact(11), Inexact(20), Absent, Absent), + (10, Inexact(0), Inexact(10), Absent, Absent), + Some(Inexact(0)), + ), + // distinct(left) = 0, distinct(right) = 0 + ( + (10, Inexact(1), Inexact(10), Inexact(0), Absent), + (10, Inexact(1), Inexact(10), Inexact(0), Absent), + None, + ), + // Inexact row count < exact null count with absent distinct count + ( + (0, Inexact(1), Inexact(10), Absent, Exact(5)), + (10, Inexact(1), Inexact(10), Absent, Absent), + Some(Inexact(0)), + ), + // NDV > num_rows: distinct count should be capped at row count + ( + (5, Inexact(1), Inexact(100), Inexact(50), Absent), + (10, Inexact(1), Inexact(100), Inexact(50), Absent), + // max_distinct_count caps: left NDV=min(50,5)=5, right NDV=min(50,10)=10 + // cardinality = (5 * 10) / max(5, 10) = 50 / 10 = 5 + Some(Inexact(5)), + ), + // NDV > num_rows on one side only + ( + (3, Inexact(1), Inexact(100), Inexact(100), Absent), + (10, Inexact(1), Inexact(100), Inexact(5), Absent), + // max_distinct_count caps: left NDV=min(100,3)=3, right NDV=min(5,10)=5 + // cardinality = (3 * 10) / max(3, 5) = 30 / 5 = 6 + Some(Inexact(6)), + ), + ]; + + for (left_info, right_info, expected_cardinality) in cases { + let left_num_rows = left_info.0; + let left_col_stats = vec![create_column_stats( + left_info.1, + left_info.2, + left_info.3, + left_info.4, + )]; + + let right_num_rows = right_info.0; + let right_col_stats = vec![create_column_stats( + right_info.1, + right_info.2, + right_info.3, + right_info.4, + )]; + + assert_eq!( + estimate_inner_join_cardinality( + Statistics { + num_rows: Inexact(left_num_rows), + total_byte_size: Absent, + column_statistics: left_col_stats.clone(), + }, + Statistics { + num_rows: Inexact(right_num_rows), + total_byte_size: Absent, + column_statistics: right_col_stats.clone(), + }, + ), + expected_cardinality.clone() + ); + + // We should also be able to use join_cardinality to get the same results + let join_type = JoinType::Inner; + let join_on = vec![( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("b", 0)) as _, + )]; + let partial_join_stats = estimate_join_cardinality( + &join_type, + create_stats(Some(left_num_rows), left_col_stats.clone(), false), + create_stats(Some(right_num_rows), right_col_stats.clone(), false), + &join_on, + NullEquality::NullEqualsNothing, + ); + + assert_eq!( + partial_join_stats.clone().map(|s| Inexact(s.num_rows)), + expected_cardinality.clone() + ); + assert_eq!( + partial_join_stats.map(|s| s.column_statistics), + expected_cardinality.map(|_| [left_col_stats, right_col_stats].concat()) + ); + } + Ok(()) + } + + #[test] + fn test_inner_join_cardinality_multiplication_overflow() { + let statistics = |num_rows, distinct_count| Statistics { + num_rows, + total_byte_size: Absent, + column_statistics: vec![ColumnStatistics { + distinct_count, + ..Default::default() + }], + }; + let large_row_count = usize::MAX / 2 + 1; + + // The Cartesian product overflows usize, but applying the NDV divisor + // produces a representable cardinality. + assert_eq!( + estimate_inner_join_cardinality( + statistics(Inexact(large_row_count), Inexact(1)), + statistics(Inexact(3), Inexact(3)), + ), + Some(Inexact(large_row_count)) + ); + assert_eq!( + estimate_inner_join_cardinality( + statistics(Exact(large_row_count), Exact(1)), + statistics(Exact(3), Exact(3)), + ), + Some(Exact(large_row_count)) + ); + + // If the normalized result itself cannot fit in usize, cap the + // estimate and mark it as inexact. + assert_eq!( + estimate_inner_join_cardinality( + statistics(Exact(usize::MAX), Exact(1)), + statistics(Exact(2), Exact(1)), + ), + Some(Inexact(usize::MAX)) + ); + } + + #[test] + fn test_inner_join_cardinality_multiple_column() -> Result<()> { + let left_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(100), Absent), + create_column_stats(Inexact(100), Inexact(500), Inexact(150), Absent), + ]; + + let right_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(50), Absent), + create_column_stats(Inexact(100), Inexact(500), Inexact(200), Absent), + ]; + + // We have statistics about 4 columns, where the highest distinct + // count is 200, so we are going to pick it. + assert_eq!( + estimate_inner_join_cardinality( + Statistics { + num_rows: Inexact(400), + total_byte_size: Absent, + column_statistics: left_col_stats, + }, + Statistics { + num_rows: Inexact(400), + total_byte_size: Absent, + column_statistics: right_col_stats, + }, + ), + Some(Inexact((400 * 400) / 200)) + ); + Ok(()) + } + + #[test] + fn test_inner_join_cardinality_decimal_range() -> Result<()> { + let left_col_stats = vec![ColumnStatistics { + distinct_count: Absent, + min_value: Inexact(ScalarValue::Decimal128(Some(32500), 14, 4)), + max_value: Inexact(ScalarValue::Decimal128(Some(35000), 14, 4)), + ..Default::default() + }]; + + let right_col_stats = vec![ColumnStatistics { + distinct_count: Absent, + min_value: Inexact(ScalarValue::Decimal128(Some(33500), 14, 4)), + max_value: Inexact(ScalarValue::Decimal128(Some(34000), 14, 4)), + ..Default::default() + }]; + + assert_eq!( + estimate_inner_join_cardinality( + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: left_col_stats, + }, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: right_col_stats, + }, + ), + Some(Inexact(100)) + ); + Ok(()) + } + + #[test] + fn test_join_cardinality() -> Result<()> { + // Left table (rows=1000) + // a: min=0, max=100, distinct=100 + // b: min=0, max=500, distinct=500 + // x: min=1000, max=10000, distinct=None + // + // Right table (rows=2000) + // c: min=0, max=100, distinct=50 + // d: min=0, max=2000, distinct=2500 (how? some inexact statistics) + // y: min=0, max=100, distinct=None + // + // Join on a=c, b=d (ignore x/y) + // Right column d has NDV=2500 but only 2000 rows, so NDV is capped + // to 2000. join_selectivity = max(500, 2000) = 2000. + // Inner cardinality = (1000 * 2000) / 2000 = 1000 + let cases = vec![ + (JoinType::Inner, 1000), + (JoinType::Left, 1000), + (JoinType::Right, 2000), + (JoinType::Full, 2000), + ]; + + let left_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(100), Absent), + create_column_stats(Inexact(0), Inexact(500), Inexact(500), Absent), + create_column_stats(Inexact(1000), Inexact(10000), Absent, Absent), + ]; + + let right_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(50), Absent), + create_column_stats(Inexact(0), Inexact(2000), Inexact(2500), Absent), + create_column_stats(Inexact(0), Inexact(100), Absent, Absent), + ]; + + for (join_type, expected_num_rows) in cases { + let join_on = vec![ + ( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("c", 0)) as _, + ), + ( + Arc::new(Column::new("b", 1)) as _, + Arc::new(Column::new("d", 1)) as _, + ), + ]; + + let partial_join_stats = estimate_join_cardinality( + &join_type, + create_stats(Some(1000), left_col_stats.clone(), false), + create_stats(Some(2000), right_col_stats.clone(), false), + &join_on, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(partial_join_stats.num_rows, expected_num_rows); + assert_eq!( + partial_join_stats.column_statistics, + [left_col_stats.clone(), right_col_stats.clone()].concat() + ); + } + + Ok(()) + } + + #[test] + fn test_join_cardinality_key_order() -> Result<()> { + // Reversing join key order should not change estimated cardinality + let left_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(100), Absent), + create_column_stats(Inexact(0), Inexact(500), Inexact(500), Absent), + create_column_stats(Inexact(1000), Inexact(10000), Absent, Absent), + ]; + + let right_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(50), Absent), + create_column_stats(Inexact(0), Inexact(2000), Inexact(2500), Absent), + create_column_stats(Inexact(0), Inexact(100), Absent, Absent), + ]; + + let join_on_ab = vec![ + ( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("c", 0)) as _, + ), + ( + Arc::new(Column::new("b", 1)) as _, + Arc::new(Column::new("d", 1)) as _, + ), + ]; + let join_on_ba = vec![ + ( + Arc::new(Column::new("b", 1)) as _, + Arc::new(Column::new("d", 1)) as _, + ), + ( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("c", 0)) as _, + ), + ]; + + let stats_ab = estimate_join_cardinality( + &JoinType::Inner, + create_stats(Some(1000), left_col_stats.clone(), false), + create_stats(Some(2000), right_col_stats.clone(), false), + &join_on_ab, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + let stats_ba = estimate_join_cardinality( + &JoinType::Inner, + create_stats(Some(1000), left_col_stats.clone(), false), + create_stats(Some(2000), right_col_stats.clone(), false), + &join_on_ba, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + assert_eq!(stats_ab.num_rows, 1000); + assert_eq!(stats_ba.num_rows, stats_ab.num_rows); + assert_eq!(stats_ba.column_statistics, stats_ab.column_statistics); + assert_eq!( + stats_ab.column_statistics, + [left_col_stats, right_col_stats].concat() + ); + + Ok(()) + } + + #[test] + fn test_join_cardinality_when_one_column_is_disjoint() -> Result<()> { + // Left table (rows=1000) + // a: min=0, max=100, distinct=100 + // b: min=0, max=500, distinct=500 + // x: min=1000, max=10000, distinct=None + // + // Right table (rows=2000) + // c: min=0, max=100, distinct=50 + // d: min=0, max=2000, distinct=2500 (how? some inexact statistics) + // y: min=0, max=100, distinct=None + // + // Join on a=c, x=y (ignores b/d) where x and y does not intersect + + let left_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(100), Absent), + create_column_stats(Inexact(0), Inexact(500), Inexact(500), Absent), + create_column_stats(Inexact(1000), Inexact(10000), Absent, Absent), + ]; + + let right_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(50), Absent), + create_column_stats(Inexact(0), Inexact(2000), Inexact(2500), Absent), + create_column_stats(Inexact(0), Inexact(100), Absent, Absent), + ]; + + let join_on = vec![ + ( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("c", 0)) as _, + ), + ( + Arc::new(Column::new("x", 2)) as _, + Arc::new(Column::new("y", 2)) as _, + ), + ]; + + let cases = vec![ + // Join type, expected cardinality + // + // When an inner join is disjoint, that means it won't + // produce any rows. + (JoinType::Inner, 0), + // But left/right outer joins will produce at least + // the amount of rows from the left/right side. + (JoinType::Left, 1000), + (JoinType::Right, 2000), + // And a full outer join will produce at least the combination + // of the rows above (minus the cardinality of the inner join, which + // is 0). + (JoinType::Full, 3000), + ]; + + for (join_type, expected_num_rows) in cases { + let partial_join_stats = estimate_join_cardinality( + &join_type, + create_stats(Some(1000), left_col_stats.clone(), true), + create_stats(Some(2000), right_col_stats.clone(), true), + &join_on, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(partial_join_stats.num_rows, expected_num_rows); + assert_eq!( + partial_join_stats.column_statistics, + [left_col_stats.clone(), right_col_stats.clone()].concat() + ); + } + + Ok(()) + } + + #[test] + fn test_anti_semi_join_cardinality() -> Result<()> { + let cases: Vec<(JoinType, PartialStats, PartialStats, Option<usize>)> = vec![ + // ------------------------------------------------ + // | join_type , | + // | left(rows, min, max, distinct, null_count), | + // | right(rows, min, max, distinct, null_count), | + // | expected, | + // ------------------------------------------------ + + // Cardinality computation + // ======================= + ( + JoinType::LeftSemi, + (50, Inexact(10), Inexact(20), Absent, Absent), + (10, Inexact(15), Inexact(25), Absent, Absent), + Some(46), + ), + ( + JoinType::RightSemi, + (50, Inexact(10), Inexact(20), Absent, Absent), + (10, Inexact(15), Inexact(25), Absent, Absent), + Some(10), + ), + ( + JoinType::LeftSemi, + (10, Absent, Absent, Absent, Absent), + (50, Absent, Absent, Absent, Absent), + Some(10), + ), + ( + JoinType::LeftSemi, + (50, Inexact(10), Inexact(20), Absent, Absent), + (10, Inexact(30), Inexact(40), Absent, Absent), + Some(0), + ), + ( + JoinType::LeftSemi, + (50, Inexact(10), Absent, Absent, Absent), + (10, Absent, Inexact(5), Absent, Absent), + Some(0), + ), + ( + JoinType::LeftSemi, + (50, Absent, Inexact(20), Absent, Absent), + (10, Inexact(30), Absent, Absent, Absent), + Some(0), + ), + ( + JoinType::LeftAnti, + (50, Inexact(10), Inexact(20), Absent, Absent), + (10, Inexact(15), Inexact(25), Absent, Absent), + Some(4), + ), + ( + JoinType::RightAnti, + (50, Inexact(10), Inexact(20), Absent, Absent), + (10, Inexact(15), Inexact(25), Absent, Absent), + Some(0), + ), + ( + JoinType::LeftAnti, + (10, Absent, Absent, Absent, Absent), + (50, Absent, Absent, Absent, Absent), + Some(10), + ), + ( + JoinType::LeftAnti, + (50, Inexact(10), Inexact(20), Absent, Absent), + (10, Inexact(30), Inexact(40), Absent, Absent), + Some(50), + ), + ( + JoinType::LeftAnti, + (50, Inexact(10), Absent, Absent, Absent), + (10, Absent, Inexact(5), Absent, Absent), + Some(50), + ), + ( + JoinType::LeftAnti, + (50, Absent, Inexact(20), Absent, Absent), + (10, Inexact(30), Absent, Absent, Absent), + Some(50), + ), + // NDV-based semi join: outer_ndv=20, inner_ndv=10 + // selectivity = 10/20 = 0.5, cardinality = ceil(50 * 0.5) = 25 + ( + JoinType::LeftSemi, + (50, Inexact(1), Inexact(100), Inexact(20), Absent), + (10, Inexact(1), Inexact(100), Inexact(10), Absent), + Some(25), + ), + // inner_ndv(30) >= outer_ndv(20) -> selectivity 1.0, no reduction + ( + JoinType::LeftSemi, + (50, Inexact(1), Inexact(100), Inexact(20), Absent), + (100, Inexact(1), Inexact(100), Inexact(30), Absent), + Some(50), + ), + // NDV-based anti join: semi=25, anti = 50 - 25 = 25 + ( + JoinType::LeftAnti, + (50, Inexact(1), Inexact(100), Inexact(20), Absent), + (10, Inexact(1), Inexact(100), Inexact(10), Absent), + Some(25), + ), + // inner covers all outer: semi=50, anti = 0 + ( + JoinType::LeftAnti, + (50, Inexact(1), Inexact(100), Inexact(20), Absent), + (100, Inexact(1), Inexact(100), Inexact(30), Absent), + Some(0), + ), + // RightSemi with explicit NDV (NDV within row count, used as-is): + // For RightSemi, sides are swapped: outer = right (20 rows, ndv=10), + // inner = left (50 rows, ndv=5). selectivity = min(10,5)/10 = 0.5, + // cardinality = ceil(20 * 0.5) = 10. + ( + JoinType::RightSemi, + (50, Inexact(1), Inexact(100), Inexact(5), Absent), + (20, Inexact(1), Inexact(100), Inexact(10), Absent), + Some(10), + ), + // RightAnti with explicit NDV: anti = outer_rows - semi = 20 - 10 = 10. + ( + JoinType::RightAnti, + (50, Inexact(1), Inexact(100), Inexact(5), Absent), + (20, Inexact(1), Inexact(100), Inexact(10), Absent), + Some(10), + ), + // RightSemi where right-side NDV (20) exceeds right-side row count (10): + // NDV is clamped to 10, so outer_ndv=10, inner_ndv=10, + // selectivity = min(10,10)/10 = 1.0, cardinality = ceil(10 * 1.0) = 10. + ( + JoinType::RightSemi, + (50, Inexact(1), Inexact(100), Inexact(10), Absent), + (10, Inexact(1), Inexact(100), Inexact(20), Absent), + Some(10), + ), + // RightAnti with NDV clamped by row count: anti = 10 - 10 = 0. + ( + JoinType::RightAnti, + (50, Inexact(1), Inexact(100), Inexact(10), Absent), + (10, Inexact(1), Inexact(100), Inexact(20), Absent), + Some(0), + ), + // Empty inner table: no match possible, semi → 0 + ( + JoinType::LeftSemi, + (100, Absent, Absent, Absent, Absent), + (0, Absent, Absent, Absent, Absent), + Some(0), + ), + // NDV-based semi with nulls on outer side: + // outer_ndv=20, inner_ndv=10, null_frac=10/100=0.1 + // selectivity = 10/20 * (1-0.1) = 0.5 * 0.9 = 0.45 + // semi = ceil(100 * 0.45) = 45 + ( + JoinType::LeftSemi, + (100, Absent, Absent, Inexact(20), Inexact(10)), + (200, Absent, Absent, Inexact(10), Absent), + Some(45), + ), + // Anti-join with nulls on outer side: + // semi=45, anti = 100 - 45 = 55 + ( + JoinType::LeftAnti, + (100, Absent, Absent, Inexact(20), Inexact(10)), + (200, Absent, Absent, Inexact(10), Absent), + Some(55), + ), + // All outer rows are null: null_frac=1.0 + // selectivity = 10/20 * (1-1.0) = 0.0, semi = 0 + ( + JoinType::LeftSemi, + (100, Absent, Absent, Inexact(20), Inexact(100)), + (200, Absent, Absent, Inexact(10), Absent), + Some(0), + ), + // All outer rows are null (anti): anti = 100 - 0 = 100 + ( + JoinType::LeftAnti, + (100, Absent, Absent, Inexact(20), Inexact(100)), + (200, Absent, Absent, Inexact(10), Absent), + Some(100), + ), + ]; + + let join_on = vec![( + Arc::new(Column::new("l_col", 0)) as _, + Arc::new(Column::new("r_col", 0)) as _, + )]; + + for (join_type, outer_info, inner_info, expected) in cases { + let outer_num_rows = outer_info.0; + let outer_col_stats = vec![create_column_stats( + outer_info.1, + outer_info.2, + outer_info.3, + outer_info.4, + )]; + + let inner_num_rows = inner_info.0; + let inner_col_stats = vec![create_column_stats( + inner_info.1, + inner_info.2, + inner_info.3, + inner_info.4, + )]; + + let output_cardinality = estimate_join_cardinality( + &join_type, + Statistics { + num_rows: Inexact(outer_num_rows), + total_byte_size: Absent, + column_statistics: outer_col_stats, + }, + Statistics { + num_rows: Inexact(inner_num_rows), + total_byte_size: Absent, + column_statistics: inner_col_stats, + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|cardinality| cardinality.num_rows); + + assert_eq!( + output_cardinality, expected, + "failure for join_type: {join_type}" + ); + } + + Ok(()) + } + + #[test] + fn test_semi_join_cardinality_absent_rows() -> Result<()> { + let dummy_column_stats = + vec![create_column_stats(Absent, Absent, Absent, Absent)]; + let join_on = vec![( + Arc::new(Column::new("l_col", 0)) as _, + Arc::new(Column::new("r_col", 0)) as _, + )]; + + let absent_outer_estimation = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Absent, + total_byte_size: Absent, + column_statistics: dummy_column_stats.clone(), + }, + Statistics { + num_rows: Exact(10), + total_byte_size: Absent, + column_statistics: dummy_column_stats.clone(), + }, + &join_on, + NullEquality::NullEqualsNothing, + ); + assert!( + absent_outer_estimation.is_none(), + "Expected \"None\" estimated SemiJoin cardinality for absent outer num_rows" + ); + + let absent_inner_estimation = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(500), + total_byte_size: Absent, + column_statistics: dummy_column_stats.clone(), + }, + Statistics { + num_rows: Absent, + total_byte_size: Absent, + column_statistics: dummy_column_stats.clone(), + }, + &join_on, + NullEquality::NullEqualsNothing, + ).expect("Expected non-empty PartialJoinStatistics for SemiJoin with absent inner num_rows"); + + assert_eq!( + absent_inner_estimation.num_rows, 500, + "Expected outer.num_rows estimated SemiJoin cardinality for absent inner num_rows" + ); + + let absent_inner_estimation = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Absent, + total_byte_size: Absent, + column_statistics: dummy_column_stats.clone(), + }, + Statistics { + num_rows: Absent, + total_byte_size: Absent, + column_statistics: dummy_column_stats, + }, + &join_on, + NullEquality::NullEqualsNothing, + ); + assert!( + absent_inner_estimation.is_none(), + "Expected \"None\" estimated SemiJoin cardinality for absent outer and inner num_rows" + ); + + Ok(()) + } + + #[test] + fn test_semi_join_multi_column_and_mixed_stats() -> Result<()> { + let join_on = vec![ + ( + Arc::new(Column::new("l_col0", 0)) as _, + Arc::new(Column::new("r_col0", 0)) as _, + ), + ( + Arc::new(Column::new("l_col1", 1)) as _, + Arc::new(Column::new("r_col1", 1)) as _, + ), + ]; + + // Multi-column: both columns have NDV on both sides. + // col0: outer_ndv=20, inner_ndv=10 → selectivity = 10/20 = 0.5 + // col1: outer_ndv=40, inner_ndv=10 → selectivity = 10/40 = 0.25 + // total selectivity = 0.5 * 0.25 = 0.125 + // semi = ceil(100 * 0.125) = 13 + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(20), Absent), + create_column_stats(Absent, Absent, Inexact(40), Absent), + ], + }, + Statistics { + num_rows: Inexact(200), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(10), Absent), + create_column_stats(Absent, Absent, Inexact(10), Absent), + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(result, Some(13), "multi-column semi join"); + + // Multi-column anti: anti = 100 - 13 = 87 + let result = estimate_join_cardinality( + &JoinType::LeftAnti, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(20), Absent), + create_column_stats(Absent, Absent, Inexact(40), Absent), + ], + }, + Statistics { + num_rows: Inexact(200), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(10), Absent), + create_column_stats(Absent, Absent, Inexact(10), Absent), + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(result, Some(87), "multi-column anti join"); + + // Mixed stats: col0 has NDV on both sides, col1 has NDV only on outer. + // col1 is skipped (either side missing), so selectivity comes from col0 only. + // col0: outer_ndv=20, inner_ndv=10 → selectivity = 0.5 + // semi = ceil(100 * 0.5) = 50 + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(20), Absent), + create_column_stats(Absent, Absent, Inexact(40), Absent), + ], + }, + Statistics { + num_rows: Inexact(200), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(10), Absent), + create_column_stats(Absent, Absent, Absent, Absent), + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(result, Some(50), "mixed stats: col1 skipped"); + + // Mixed stats: neither column has stats on both sides → fallback to outer_rows + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(20), Absent), + create_column_stats(Absent, Absent, Absent, Absent), + ], + }, + Statistics { + num_rows: Inexact(200), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Absent, Absent), + create_column_stats(Absent, Absent, Inexact(10), Absent), + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(result, Some(100), "no column has stats on both sides"); + + // Multi-column with nulls on one column: + // col0: outer_ndv=20, inner_ndv=10, null_frac=0.0 → 10/20 * 1.0 = 0.5 + // col1: outer_ndv=40, inner_ndv=10, null_frac=20/100=0.2 → 10/40 * 0.8 = 0.2 + // total selectivity = 0.5 * 0.2 = 0.1 + // semi = ceil(100 * 0.1) = 10 + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(20), Absent), + create_column_stats(Absent, Absent, Inexact(40), Inexact(20)), + ], + }, + Statistics { + num_rows: Inexact(200), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Absent, Absent, Inexact(10), Absent), + create_column_stats(Absent, Absent, Inexact(10), Absent), + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!( + result, + Some(10), + "multi-column semi join with nulls on one column" + ); + + Ok(()) + } + + #[test] + fn test_semi_anti_join_disjoint_check_uses_only_join_keys() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + // Ranges for the join key overlap; ranges for the other column are disjoint + let left_stats = Statistics { + num_rows: Inexact(50), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Inexact(1), Inexact(10), Absent, Absent), + create_column_stats(Inexact(100), Inexact(200), Absent, Absent), + ], + }; + let right_stats = Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Inexact(1), Inexact(10), Absent, Absent), + create_column_stats(Inexact(1000), Inexact(2000), Absent, Absent), + ], + }; + + let left_semi = estimate_join_cardinality( + &JoinType::LeftSemi, + left_stats.clone(), + right_stats.clone(), + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(left_semi, Some(50)); + + let left_anti = estimate_join_cardinality( + &JoinType::LeftAnti, + left_stats, + right_stats, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(left_anti, Some(0)); + } + + #[test] + fn test_semi_join_scales_preserved_column_statistics() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(432_187), + total_byte_size: Absent, + column_statistics: vec![ + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Absent, + distinct_count: Absent, + byte_size: Exact(3_457_496), + }, + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Exact(ScalarValue::from(1_000_000_i64)), + distinct_count: Exact(500_000), + byte_size: Exact(3_457_496), + }, + ], + }, + Statistics { + num_rows: Inexact(32), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Inexact(1), + Inexact(32), + Absent, + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 32); + assert_eq!(result.total_byte_size, Inexact(512)); + assert_eq!(result.column_statistics[0].null_count, Exact(0)); + assert_eq!(result.column_statistics[0].distinct_count, Absent); + assert_eq!( + result.column_statistics[0].min_value, + Inexact(ScalarValue::from(1_i64)) + ); + assert_eq!( + result.column_statistics[0].max_value, + Inexact(ScalarValue::from(432_187_i64)) + ); + assert_eq!(result.column_statistics[0].byte_size, Inexact(256)); + assert_eq!(result.column_statistics[1].null_count, Inexact(1)); + // distinct_count is capped at the non-null output rows (32 - 1). + assert_eq!(result.column_statistics[1].distinct_count, Inexact(31)); + assert_eq!(result.column_statistics[1].sum_value, Absent); + assert_eq!(result.column_statistics[1].byte_size, Inexact(256)); + } + + #[test] + fn test_semi_join_null_equals_null_scales_join_key_nulls() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(100), + Exact(20), + )], + }, + Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(10), + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNull, + ) + .expect("semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 10); + assert_eq!(result.column_statistics[0].null_count, Inexact(2)); + assert_eq!(result.column_statistics[0].distinct_count, Inexact(8)); + } + + #[test] + fn test_semi_join_total_byte_size_absent_if_any_column_byte_size_absent() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + ColumnStatistics { + null_count: Exact(0), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(100_i64)), + sum_value: Absent, + distinct_count: Absent, + byte_size: Exact(800), + }, + ColumnStatistics { + null_count: Exact(0), + min_value: Absent, + max_value: Absent, + sum_value: Absent, + distinct_count: Absent, + byte_size: Absent, + }, + ], + }, + Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Inexact(1), + Inexact(10), + Absent, + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 10); + assert_eq!(result.total_byte_size, Absent); + } + + #[test] + fn test_anti_join_preserves_join_key_nulls() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftAnti, + Statistics { + num_rows: Inexact(1_000_000), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(900_000), + Exact(100_000), + )], + }, + Statistics { + num_rows: Inexact(900_000), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(900_000), + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("anti join cardinality should be estimated"); + + assert_eq!(result.num_rows, 100_000); + assert_eq!(result.column_statistics[0].null_count, Inexact(100_000)); + assert_eq!(result.column_statistics[0].distinct_count, Inexact(0)); + } + + #[test] + fn test_anti_join_null_equals_null_scales_join_key_nulls() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftAnti, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(100), + Exact(20), + )], + }, + Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(10), + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNull, + ) + .expect("anti join cardinality should be estimated"); + + assert_eq!(result.num_rows, 90); + assert_eq!(result.column_statistics[0].null_count, Inexact(18)); + assert_eq!(result.column_statistics[0].distinct_count, Inexact(72)); + } + + #[test] + fn test_right_semi_join_scales_preserved_column_statistics() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + // For a right semi join the right input is preserved, so its column + // statistics (and right join-key index) are the ones normalized. + let result = estimate_join_cardinality( + &JoinType::RightSemi, + Statistics { + num_rows: Inexact(32), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Inexact(1), + Inexact(32), + Absent, + Absent, + )], + }, + Statistics { + num_rows: Inexact(432_187), + total_byte_size: Absent, + column_statistics: vec![ + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Absent, + distinct_count: Absent, + byte_size: Exact(3_457_496), + }, + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Exact(ScalarValue::from(1_000_000_i64)), + distinct_count: Exact(500_000), + byte_size: Exact(3_457_496), + }, + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("right semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 32); + // Join-key column: null counts collapse to exact zero (null keys never match). + assert_eq!(result.column_statistics[0].null_count, Exact(0)); + assert_eq!(result.column_statistics[0].byte_size, Inexact(256)); + // Non-key column: counts scaled to the subset, sum dropped, distinct + // capped at the non-null output rows (32 - 1). + assert_eq!(result.column_statistics[1].null_count, Inexact(1)); + assert_eq!(result.column_statistics[1].distinct_count, Inexact(31)); + assert_eq!(result.column_statistics[1].sum_value, Absent); + assert_eq!(result.column_statistics[1].byte_size, Inexact(256)); + } + + #[test] + fn test_adjust_right_output_partitioning_preserves_range() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(Some(100)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(20)), + ScalarValue::Int32(Some(50)), + ]), + ]; + let range = RangePartitioning::try_new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 2)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points.clone(), + )?; + + let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?; + let expected = Partitioning::Range(RangePartitioning::new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 3)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 5)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points, + )); + + assert_eq!(adjusted, expected); + Ok(()) + } + + #[test] + fn test_calculate_join_output_ordering() -> Result<()> { + let left_ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))), + PhysicalSortExpr::new_default(Arc::new(Column::new("c", 2))), + PhysicalSortExpr::new_default(Arc::new(Column::new("d", 3))), + ]); + let right_ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new_default(Arc::new(Column::new("z", 2))), + PhysicalSortExpr::new_default(Arc::new(Column::new("y", 1))), + ]); + let join_type = JoinType::Inner; + let left_columns_len = 5; + let maintains_input_orders = [[true, false], [false, true]]; + let probe_sides = [Some(JoinSide::Left), Some(JoinSide::Right)]; + + let expected = [ + LexOrdering::new(vec![ + PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))), + PhysicalSortExpr::new_default(Arc::new(Column::new("c", 2))), + PhysicalSortExpr::new_default(Arc::new(Column::new("d", 3))), + PhysicalSortExpr::new_default(Arc::new(Column::new("z", 7))), + PhysicalSortExpr::new_default(Arc::new(Column::new("y", 6))), + ]), + LexOrdering::new(vec![ + PhysicalSortExpr::new_default(Arc::new(Column::new("z", 7))), + PhysicalSortExpr::new_default(Arc::new(Column::new("y", 6))), + PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))), + PhysicalSortExpr::new_default(Arc::new(Column::new("c", 2))), + PhysicalSortExpr::new_default(Arc::new(Column::new("d", 3))), + ]), + ]; + + for (i, (maintains_input_order, probe_side)) in + maintains_input_orders.iter().zip(probe_sides).enumerate() + { + assert_eq!( + calculate_join_output_ordering( + left_ordering.as_ref(), + right_ordering.as_ref(), + join_type, + left_columns_len, + maintains_input_order, + probe_side, + )?, + expected[i] + ); + } + + Ok(()) + } + + fn create_test_batch(num_rows: usize) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let data = Arc::new(Int32Array::from_iter_values(0..num_rows as i32)); + RecordBatch::try_new(schema, vec![data]).unwrap() + } + + fn assert_split_batches( + batches: Vec<(RecordBatch, bool)>, + batch_size: usize, + num_rows: usize, + ) { + let mut row_count = 0; + for (batch, last) in batches.into_iter() { + assert_eq!(batch.num_rows(), (num_rows - row_count).min(batch_size)); + let column = batch + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + for i in 0..batch.num_rows() { + assert_eq!(column.value(i), i as i32 + row_count as i32); + } + row_count += batch.num_rows(); + assert_eq!(last, row_count == num_rows); + } + } + + #[rstest] + #[test] + fn test_batch_splitter( + #[values(1, 3, 11)] batch_size: usize, + #[values(1, 6, 50)] num_rows: usize, + ) { + let mut splitter = BatchSplitter::new(batch_size); + splitter.set_batch(create_test_batch(num_rows)); + + let mut batches = Vec::with_capacity(num_rows.div_ceil(batch_size)); + while let Some(batch) = splitter.next() { + batches.push(batch); + } + + assert!(splitter.next().is_none()); + assert_split_batches(batches, batch_size, num_rows); + } + + #[tokio::test] + async fn test_swap_reverting_projection() { + let left_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + + let right_schema = Schema::new(vec![Field::new("c", DataType::Int32, false)]); + + let proj = swap_reverting_projection(&left_schema, &right_schema); + + assert_eq!(proj.len(), 3); + + let proj_expr = &proj[0]; + assert_eq!(proj_expr.alias, "a"); + assert_col_expr(&proj_expr.expr, "a", 1); + + let proj_expr = &proj[1]; + assert_eq!(proj_expr.alias, "b"); + assert_col_expr(&proj_expr.expr, "b", 2); + + let proj_expr = &proj[2]; + assert_eq!(proj_expr.alias, "c"); + assert_col_expr(&proj_expr.expr, "c", 0); + } + + fn assert_col_expr(expr: &Arc<dyn PhysicalExpr>, name: &str, index: usize) { + let col = expr + .downcast_ref::<Column>() + .expect("Projection items should be Column expression"); + assert_eq!(col.name(), name); + assert_eq!(col.index(), index); + } + + #[test] + fn test_join_metadata() -> Result<()> { + let left_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]) + .with_metadata(HashMap::from([("key".to_string(), "left".to_string())])); + + let right_schema = Schema::new(vec![Field::new("b", DataType::Int32, false)]) + .with_metadata(HashMap::from([("key".to_string(), "right".to_string())])); + + let (join_schema, _) = + build_join_schema(&left_schema, &right_schema, &JoinType::Left); + assert_eq!( + join_schema.metadata(), + &HashMap::from([("key".to_string(), "left".to_string())]) + ); + let (join_schema, _) = + build_join_schema(&left_schema, &right_schema, &JoinType::Right); + assert_eq!( + join_schema.metadata(), + &HashMap::from([("key".to_string(), "right".to_string())]) + ); + + Ok(()) + } + + #[test] + fn test_build_batch_empty_build_side_empty_schema() -> Result<()> { + // When the output schema has no fields (empty projection pushed into + // the join), build_batch_empty_build_side should return a RecordBatch + // with the correct row count but no columns. + let empty_schema = Schema::empty(); + + let build_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + + let probe_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])), + vec![Arc::new(Int32Array::from(vec![4, 5, 6, 7]))], + )?; + + let result = build_batch_empty_build_side( + &empty_schema, + &build_batch, + &probe_batch, + &[], // no column indices with empty projection + JoinType::Right, + )?; + + assert_eq!(result.num_rows(), 4); + assert_eq!(result.num_columns(), 0); + + Ok(()) + } + + #[test] + fn test_max_distinct_count_no_overflow_when_null_count_exceeds_num_rows() { + let num_rows = Exact(2); + let stats = ColumnStatistics { + distinct_count: Absent, + null_count: Exact(5), + min_value: Absent, + max_value: Absent, + sum_value: Absent, + byte_size: Absent, + }; + let result = max_distinct_count(&num_rows, &stats); + assert_eq!(result, Exact(0)); + } + + #[test] + fn test_join_key_comparator_multi_column() { + let left_a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 2, 3])); + let left_b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + let right_a: ArrayRef = Arc::new(Int32Array::from(vec![2, 2, 3, 4])); + let right_b: ArrayRef = Arc::new(StringArray::from(vec!["b", "d", "a", "a"])); + + let opts = vec![SortOptions::default(), SortOptions::default()]; + let cmp = JoinKeyComparator::new( + &[left_a, left_b], + &[right_a, right_b], + &opts, + NullEquality::NullEqualsNull, + ) + .unwrap(); + + // left[0]=(1,"a") vs right[0]=(2,"b") -> Less (first column) + assert_eq!(cmp.compare(0, 0), Ordering::Less); + // left[1]=(2,"b") vs right[0]=(2,"b") -> Equal + assert_eq!(cmp.compare(1, 0), Ordering::Equal); + assert!(cmp.is_equal(1, 0)); + // left[2]=(2,"c") vs right[1]=(2,"d") -> Less (second column) + assert_eq!(cmp.compare(2, 1), Ordering::Less); + // left[3]=(3,"d") vs right[0]=(2,"b") -> Greater + assert_eq!(cmp.compare(3, 0), Ordering::Greater); + } + + #[test] + fn test_join_key_comparator_null_equals_null() { + let left: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), None, None, Some(2)])); + let right: ArrayRef = + Arc::new(Int32Array::from(vec![None, None, Some(1), Some(2)])); + + let opts = vec![SortOptions { + descending: false, + nulls_first: true, + }]; + let cmp = JoinKeyComparator::new( + &[left], + &[right], + &opts, + NullEquality::NullEqualsNull, + ) + .unwrap(); + + // left[1]=NULL vs right[1]=NULL -> Equal (NullEqualsNull) + assert_eq!(cmp.compare(1, 1), Ordering::Equal); + assert!(cmp.is_equal(1, 1)); + // left[0]=1 vs right[0]=NULL -> Greater (nulls_first, non-null > null) + assert_eq!(cmp.compare(0, 0), Ordering::Greater); + // left[3]=2 vs right[3]=2 -> Equal + assert_eq!(cmp.compare(3, 3), Ordering::Equal); + } + + #[test] + fn test_join_key_comparator_null_equals_nothing() { + let left: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), None, None, Some(2)])); + let right: ArrayRef = + Arc::new(Int32Array::from(vec![None, None, Some(1), Some(2)])); + + let opts = vec![SortOptions { + descending: false, + nulls_first: true, + }]; + let cmp = JoinKeyComparator::new( + &[left], + &[right], + &opts, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + // left[1]=NULL vs right[1]=NULL -> Less (NullEqualsNothing) + assert_eq!(cmp.compare(1, 1), Ordering::Less); + // left[0]=1 vs right[0]=NULL -> Greater (nulls_first) + assert_eq!(cmp.compare(0, 0), Ordering::Greater); + // left[3]=2 vs right[3]=2 -> Equal + assert_eq!(cmp.compare(3, 3), Ordering::Equal); + } + + #[test] + fn test_join_key_comparator_nulls_first_ordering() { + let left: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(1)])); + let right: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None])); + + // nulls_first = true: null < non-null + let cmp_nf = JoinKeyComparator::new( + &[Arc::clone(&left)], + &[Arc::clone(&right)], + &[SortOptions { + descending: false, + nulls_first: true, + }], + NullEquality::NullEqualsNull, + ) + .unwrap(); + assert_eq!(cmp_nf.compare(0, 0), Ordering::Less); + assert_eq!(cmp_nf.compare(1, 1), Ordering::Greater); + + // nulls_first = false: null > non-null + let cmp_nl = JoinKeyComparator::new( + &[left], + &[right], + &[SortOptions { + descending: false, + nulls_first: false, + }], + NullEquality::NullEqualsNull, + ) + .unwrap(); + assert_eq!(cmp_nl.compare(0, 0), Ordering::Greater); + assert_eq!(cmp_nl.compare(1, 1), Ordering::Less); + } + + #[test] + fn test_equal_rows_arr_filters_candidate_pairs() { + let left_a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 2, 3])); + let left_b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + let right_a: ArrayRef = Arc::new(Int32Array::from(vec![2, 2, 3, 4])); + let right_b: ArrayRef = Arc::new(StringArray::from(vec!["b", "d", "d", "a"])); + + let left_indices = UInt64Array::from(vec![0, 1, 2, 3]); + let right_indices = UInt32Array::from(vec![0, 0, 1, 2]); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[left_a, left_b], + &[right_a, right_b], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + assert_eq!(left_filtered, UInt64Array::from(vec![1, 3])); + assert_eq!(right_filtered, UInt32Array::from(vec![0, 2])); + } + + #[test] + fn test_equal_rows_arr_empty_keys_returns_empty() { + let left_indices = UInt64Array::from(vec![0, 1, 2]); + let right_indices = UInt32Array::from(vec![0, 1, 2]); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[], + &[], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + assert_eq!(left_filtered.len(), 0); + assert_eq!(right_filtered.len(), 0); + } + + #[test] + fn test_equal_rows_arr_respects_null_equality() { + let left: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None])); + let right: ArrayRef = + Arc::new(Int32Array::from(vec![None, Some(1), Some(2), None])); + let left_indices = UInt64Array::from(vec![0, 1, 2, 3]); + let right_indices = UInt32Array::from(vec![1, 0, 2, 3]); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[Arc::clone(&left)], + &[Arc::clone(&right)], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0, 2])); + assert_eq!(right_filtered, UInt32Array::from(vec![1, 2])); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[left], + &[right], + NullEquality::NullEqualsNull, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0, 1, 2, 3])); + assert_eq!(right_filtered, UInt32Array::from(vec![1, 0, 2, 3])); + } + + #[test] + fn test_equal_rows_arr_single_string_col_fast_path() { + // Single-column string keys exercise the specialized fast path, + // including null handling under both null-equality modes. + let left: ArrayRef = Arc::new(StringArray::from(vec![ + Some("long_shared_join_key_value"), + None, + Some("long_shared_join_key_value"), + Some("other"), + ])); + let right: ArrayRef = Arc::new(StringArray::from(vec![ + Some("long_shared_join_key_value"), + None, + Some("mismatch"), + None, + ])); + let left_indices = UInt64Array::from(vec![0, 1, 2, 3]); + let right_indices = UInt32Array::from(vec![0, 1, 2, 3]); + + // NullEqualsNothing: only the (0,0) value pair matches; both-null drops. + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[Arc::clone(&left)], + &[Arc::clone(&right)], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0])); + assert_eq!(right_filtered, UInt32Array::from(vec![0])); + + // NullEqualsNull: the both-null (1,1) pair now also matches. + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[left], + &[right], + NullEquality::NullEqualsNull, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0, 1])); + assert_eq!(right_filtered, UInt32Array::from(vec![0, 1])); + } + + #[test] + fn test_equal_rows_arr_single_col_covers_all_specialized_types() { + // Drive every specialized single-column fast-path arm. Each case has a + // matching pair at index 0 and a non-matching pair at index 1, so a + // correct arm keeps exactly the first pair. + fn check(left: ArrayRef, right: ArrayRef) { + let (left_filtered, right_filtered) = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0, 1]), + &[left], + &[right], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0])); + assert_eq!(right_filtered, UInt32Array::from(vec![0])); + } + + check( + Arc::new(BooleanArray::from(vec![true, false])), + Arc::new(BooleanArray::from(vec![true, true])), + ); + check( + Arc::new(Int8Array::from(vec![1, 2])), + Arc::new(Int8Array::from(vec![1, 3])), + ); + check( + Arc::new(Int16Array::from(vec![1, 2])), + Arc::new(Int16Array::from(vec![1, 3])), + ); + check( + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(Int64Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt8Array::from(vec![1, 2])), + Arc::new(UInt8Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt16Array::from(vec![1, 2])), + Arc::new(UInt16Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt32Array::from(vec![1, 2])), + Arc::new(UInt32Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt64Array::from(vec![1, 2])), + Arc::new(UInt64Array::from(vec![1, 3])), + ); + check( + Arc::new(Decimal128Array::from(vec![1i128, 2])), + Arc::new(Decimal128Array::from(vec![1i128, 3])), + ); + check( + Arc::new(BinaryArray::from_iter_values([b"a".as_ref(), b"b"])), + Arc::new(BinaryArray::from_iter_values([b"a".as_ref(), b"c"])), + ); + check( + Arc::new(LargeBinaryArray::from_iter_values([b"a".as_ref(), b"b"])), + Arc::new(LargeBinaryArray::from_iter_values([b"a".as_ref(), b"c"])), + ); + check( + Arc::new(BinaryViewArray::from_iter_values([b"a".as_ref(), b"b"])), + Arc::new(BinaryViewArray::from_iter_values([b"a".as_ref(), b"c"])), + ); + check( + Arc::new( + FixedSizeBinaryArray::try_from_iter([[1u8], [2u8]].into_iter()).unwrap(), + ), + Arc::new( + FixedSizeBinaryArray::try_from_iter([[1u8], [3u8]].into_iter()).unwrap(), + ), + ); + check( + Arc::new(LargeStringArray::from(vec!["a", "b"])), + Arc::new(LargeStringArray::from(vec!["a", "c"])), + ); + check( + Arc::new(StringViewArray::from(vec!["a", "b"])), + Arc::new(StringViewArray::from(vec!["a", "c"])), + ); + check( + Arc::new(Date32Array::from(vec![1, 2])), + Arc::new(Date32Array::from(vec![1, 3])), + ); + check( + Arc::new(Date64Array::from(vec![1, 2])), + Arc::new(Date64Array::from(vec![1, 3])), + ); + check( + Arc::new(TimestampSecondArray::from(vec![1, 2])), + Arc::new(TimestampSecondArray::from(vec![1, 3])), + ); + check( + Arc::new(TimestampMillisecondArray::from(vec![1, 2])), + Arc::new(TimestampMillisecondArray::from(vec![1, 3])), + ); + check( + Arc::new(TimestampMicrosecondArray::from(vec![1, 2])), + Arc::new(TimestampMicrosecondArray::from(vec![1, 3])), + ); + check( + Arc::new(TimestampNanosecondArray::from(vec![1, 2])), + Arc::new(TimestampNanosecondArray::from(vec![1, 3])), + ); + } + + #[test] + fn test_equal_rows_arr_single_float_col_uses_general_path() { + // Floats are intentionally not specialized: the fast path returns + // `None` and the general comparator handles them (covers the + // fall-through arm). + let left: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + let right: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 3.0])); + let (left_filtered, right_filtered) = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0, 1]), + &[left], + &[right], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0])); + assert_eq!(right_filtered, UInt32Array::from(vec![0])); + } + + #[test] + fn test_equal_rows_arr_rejects_mismatched_inputs() { + let left: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let right: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + + let err = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0]), + &[Arc::clone(&left)], + &[Arc::clone(&right)], + NullEquality::NullEqualsNothing, + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("Cannot compare join indices with different lengths") + ); + + let err = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0, 1]), + &[left, Arc::new(Int32Array::from(vec![3, 4]))], + &[right], + NullEquality::NullEqualsNothing, + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("Cannot compare join keys with different column counts") + ); + } + + #[test] + fn test_max_distinct_count_preserves_precision_when_not_capped() { + assert_eq!( + max_distinct_count( + &Exact(10), + &ColumnStatistics { + distinct_count: Exact(5), + ..Default::default() + } + ), + Exact(5) + ); + assert_eq!( + max_distinct_count( + &Exact(10), + &ColumnStatistics { + distinct_count: Inexact(5), + ..Default::default() + } + ), + Inexact(5) + ); + // Inexact num_rows does not affect an exact NDV that is within bounds + assert_eq!( + max_distinct_count( + &Inexact(10), + &ColumnStatistics { + distinct_count: Exact(5), + ..Default::default() + } + ), + Exact(5) + ); + } + + #[test] + fn test_max_distinct_count_demotes_to_inexact_when_capped() { + // Exact NDV > Exact num_rows is an illegal state (NDV <= num_rows is a + // mathematical invariant), but the code handles it defensively by + // capping and demoting to inexact + assert_eq!( + max_distinct_count( + &Exact(10), + &ColumnStatistics { + distinct_count: Exact(15), + ..Default::default() + } + ), + Inexact(10) + ); + assert_eq!( + max_distinct_count( + &Inexact(10), + &ColumnStatistics { + distinct_count: Exact(15), + ..Default::default() + } + ), + Inexact(10) + ); + assert_eq!( + max_distinct_count( + &Exact(10), + &ColumnStatistics { + distinct_count: Inexact(15), + ..Default::default() + } + ), + Inexact(10) + ); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/lib.rs b/native/vendor/datafusion-physical-plan/src/lib.rs new file mode 100644 index 0000000..9e50a93 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/lib.rs
@@ -0,0 +1,115 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_cfg))] +// Make sure fast / cheap clones on Arc are explicit: +// https://github.com/apache/datafusion/issues/11143 +#![deny(clippy::clone_on_ref_ptr)] +#![cfg_attr(test, allow(clippy::needless_pass_by_value))] + +//! Traits for physical query plan, supporting parallel execution for partitioned relations. +//! +//! Entrypoint of this crate is trait [ExecutionPlan]. + +pub use datafusion_common::hash_utils; +pub use datafusion_common::utils::project_schema; +pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; +pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +pub use datafusion_expr::{Accumulator, ColumnarValue}; +use datafusion_physical_expr::PhysicalSortExpr; +pub use datafusion_physical_expr::window::WindowExpr; +pub use datafusion_physical_expr::{ + Distribution, Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, expressions, +}; + +pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +pub use crate::distribution_requirements::{ + ChildSatisfactionOptions, InputDistributionRequirements, +}; +#[expect(deprecated)] +pub use crate::execution_plan::{ + AsPhysicalExprRef, ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, apply_expression_roots, collect, + collect_partitioned, displayable, execute_input_stream, execute_stream, + execute_stream_partitioned, get_plan_string, replace_children_if_necessary, + with_new_children_if_necessary, +}; +pub use crate::metrics::Metric; +pub use crate::ordering::InputOrderMode; +pub use crate::sort_pushdown::SortOrderPushdownResult; +pub use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; +pub use crate::stream::EmptyRecordBatchStream; +pub use crate::topk::TopK; +pub use crate::visitor::{ExecutionPlanVisitor, accept, visit_execution_plan}; +pub use crate::work_table::WorkTable; +pub use spill::spill_manager::SpillManager; + +mod ordering; +mod render_tree; +mod topk; +mod visitor; + +pub mod aggregates; +pub mod analyze; +pub mod async_func; +pub mod buffer; +pub mod coalesce; +pub mod coalesce_batches; +pub mod coalesce_partitions; +pub mod column_rewriter; +pub mod common; +pub mod coop; +pub mod display; +pub mod distribution_requirements; +pub mod empty; +pub mod execution_plan; +pub mod explain; +pub mod filter; +pub mod filter_pushdown; +pub mod joins; +pub mod limit; +pub mod memory; +pub mod metrics; +pub mod operator_statistics; +pub mod placeholder_row; +pub mod projection; +#[cfg(feature = "proto")] +pub mod proto; +pub mod recursive_query; +pub mod repartition; +pub mod scalar_subquery; +pub mod sort_pushdown; +pub mod sorts; +pub mod spill; +pub mod statistics; +pub mod stream; +pub mod streaming; +pub mod tree_node; +pub mod union; +pub mod unnest; +pub mod windows; +pub mod work_table; +pub mod udaf { + pub use datafusion_expr::StatisticsArgs; + pub use datafusion_physical_expr::aggregate::AggregateFunctionExpr; +} + +pub mod test;
diff --git a/native/vendor/datafusion-physical-plan/src/limit.rs b/native/vendor/datafusion-physical-plan/src/limit.rs new file mode 100644 index 0000000..dd62c93 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/limit.rs
@@ -0,0 +1,1094 @@ +// 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. + +//! Defines the LIMIT plan + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use super::{ + DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, Statistics, +}; +use crate::execution_plan::{Boundedness, CardinalityEffect}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, validate_child_count, +}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_execution::TaskContext; + +use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; +use futures::stream::{Stream, StreamExt}; +use log::trace; + +/// Limit execution plan +#[derive(Debug, Clone)] +pub struct GlobalLimitExec { + /// Input execution plan + input: Arc<dyn ExecutionPlan>, + /// Number of rows to skip before fetch + skip: usize, + /// Maximum number of rows to fetch, + /// `None` means fetching all rows + fetch: Option<usize>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Input ordering that must be preserved so limit pushdown does not change + /// which rows are returned. + required_ordering: Option<LexOrdering>, + cache: Arc<PlanProperties>, +} + +impl GlobalLimitExec { + /// Create a new GlobalLimitExec + pub fn new(input: Arc<dyn ExecutionPlan>, skip: usize, fetch: Option<usize>) -> Self { + let cache = Self::compute_properties(&input); + GlobalLimitExec { + input, + skip, + fetch, + metrics: ExecutionPlanMetricsSet::new(), + required_ordering: None, + cache: Arc::new(cache), + } + } + + /// Input execution plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Number of rows to skip before fetch + pub fn skip(&self) -> usize { + self.skip + } + + /// Maximum number of rows to fetch + pub fn fetch(&self) -> Option<usize> { + self.fetch + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties { + PlanProperties::new( + input.equivalence_properties().clone(), // Equivalence Properties + Partitioning::UnknownPartitioning(1), // Output Partitioning + input.pipeline_behavior(), + // Limit operations are always bounded since they output a finite number of rows + Boundedness::Bounded, + ) + } + + /// Get the required ordering from limit + pub fn required_ordering(&self) -> &Option<LexOrdering> { + &self.required_ordering + } + + /// Set the required ordering for limit + pub fn set_required_ordering(&mut self, required_ordering: Option<LexOrdering>) { + self.required_ordering = required_ordering; + } +} + +impl DisplayAs for GlobalLimitExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "GlobalLimitExec: skip={}, fetch={}", + self.skip, + self.fetch + .map_or_else(|| "None".to_string(), |x| x.to_string()) + ) + } + DisplayFormatType::TreeRender => { + if let Some(fetch) = self.fetch { + writeln!(f, "limit={fetch}")?; + } + write!(f, "skip={}", self.skip) + } + } + } +} + +impl ExecutionPlan for GlobalLimitExec { + fn name(&self) -> &'static str { + "GlobalLimitExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut new_limit = + GlobalLimitExec::new(children.swap_remove(0), self.skip, self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!("Start GlobalLimitExec::execute for partition: {partition}"); + // GlobalLimitExec has a single output partition + assert_eq_or_internal_err!( + partition, + 0, + "GlobalLimitExec invalid partition {partition}" + ); + + // GlobalLimitExec requires a single input partition + assert_eq_or_internal_err!( + self.input.output_partitioning().partition_count(), + 1, + "GlobalLimitExec requires a single input partition" + ); + + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + let stream = self.input.execute(0, context)?; + Ok(Box::pin(LimitStream::new( + stream, + self.skip, + self.fetch, + baseline_metrics, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, self.skip, 1)?)) + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto; + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let required_ordering = optional_ordering_try_to_proto( + self.required_ordering.as_ref(), + &ctx.expr_ctx(), + )?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new( + protobuf::GlobalLimitExecNode { + input: Some(Box::new(input)), + skip: self.skip() as u32, + fetch: match self.fetch() { + Some(n) => n as i64, + _ => -1, // no limit + }, + required_ordering, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl GlobalLimitExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto; + use datafusion_proto_models::protobuf; + let limit = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit, + "GlobalLimitExec", + ); + let input = ctx.decode_required_child( + limit.input.as_deref(), + "GlobalLimitExec", + "input", + )?; + let fetch = if limit.fetch >= 0 { + Some(limit.fetch as usize) + } else { + None + }; + let required_ordering = optional_ordering_try_from_proto( + &limit.required_ordering, + &ctx.expr_ctx(input.schema().as_ref()), + )?; + let mut exec = GlobalLimitExec::new(input, limit.skip as usize, fetch); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) + } +} + +/// LocalLimitExec applies a limit to a single partition +#[derive(Debug, Clone)] +pub struct LocalLimitExec { + /// Input execution plan + input: Arc<dyn ExecutionPlan>, + /// Maximum number of rows to return + fetch: usize, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Input ordering that must be preserved so limit pushdown does not change + /// which rows are returned. + required_ordering: Option<LexOrdering>, + cache: Arc<PlanProperties>, +} + +impl LocalLimitExec { + /// Create a new LocalLimitExec partition + pub fn new(input: Arc<dyn ExecutionPlan>, fetch: usize) -> Self { + let cache = Self::compute_properties(&input); + Self { + input, + fetch, + metrics: ExecutionPlanMetricsSet::new(), + required_ordering: None, + cache: Arc::new(cache), + } + } + + /// Input execution plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Maximum number of rows to fetch + pub fn fetch(&self) -> usize { + self.fetch + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties { + PlanProperties::new( + input.equivalence_properties().clone(), // Equivalence Properties + input.output_partitioning().clone(), // Output Partitioning + input.pipeline_behavior(), + // Limit operations are always bounded since they output a finite number of rows + Boundedness::Bounded, + ) + } + + /// Get the required ordering from limit + pub fn required_ordering(&self) -> &Option<LexOrdering> { + &self.required_ordering + } + + /// Set the required ordering for limit + pub fn set_required_ordering(&mut self, required_ordering: Option<LexOrdering>) { + self.required_ordering = required_ordering; + } +} + +impl DisplayAs for LocalLimitExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "LocalLimitExec: fetch={}", self.fetch) + } + DisplayFormatType::TreeRender => { + write!(f, "limit={}", self.fetch) + } + } + } +} + +impl ExecutionPlan for LocalLimitExec { + fn name(&self) -> &'static str { + "LocalLimitExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut new_limit = + LocalLimitExec::new(children.swap_remove(0), self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start LocalLimitExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + let stream = self.input.execute(partition, context)?; + Ok(Box::pin(LimitStream::new( + stream, + 0, + Some(self.fetch), + baseline_metrics, + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(Some(self.fetch), 0, 1)?)) + } + + fn fetch(&self) -> Option<usize> { + Some(self.fetch) + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto; + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let required_ordering = optional_ordering_try_to_proto( + self.required_ordering.as_ref(), + &ctx.expr_ctx(), + )?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new( + protobuf::LocalLimitExecNode { + input: Some(Box::new(input)), + fetch: self.fetch() as u32, + required_ordering, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl LocalLimitExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto; + use datafusion_proto_models::protobuf; + let limit = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::LocalLimit, + "LocalLimitExec", + ); + let input = + ctx.decode_required_child(limit.input.as_deref(), "LocalLimitExec", "input")?; + let required_ordering = optional_ordering_try_from_proto( + &limit.required_ordering, + &ctx.expr_ctx(input.schema().as_ref()), + )?; + let mut exec = LocalLimitExec::new(input, limit.fetch as usize); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) + } +} + +/// A Limit stream skips `skip` rows, and then fetch up to `fetch` rows. +pub struct LimitStream { + /// The remaining number of rows to skip + skip: usize, + /// The remaining number of rows to produce + fetch: usize, + /// The input to read from. This is set to None once the limit is + /// reached to enable early termination + input: Option<SendableRecordBatchStream>, + /// Copy of the input schema + schema: SchemaRef, + /// Execution time metrics + baseline_metrics: BaselineMetrics, +} + +impl LimitStream { + pub fn new( + input: SendableRecordBatchStream, + skip: usize, + fetch: Option<usize>, + baseline_metrics: BaselineMetrics, + ) -> Self { + let schema = input.schema(); + Self { + skip, + fetch: fetch.unwrap_or(usize::MAX), + input: Some(input), + schema, + baseline_metrics, + } + } + + fn poll_and_skip( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + let input = self.input.as_mut().unwrap(); + loop { + let poll = input.poll_next_unpin(cx); + let poll = poll.map_ok(|batch| { + if batch.num_rows() <= self.skip { + self.skip -= batch.num_rows(); + RecordBatch::new_empty(input.schema()) + } else { + let new_batch = batch.slice(self.skip, batch.num_rows() - self.skip); + self.skip = 0; + new_batch + } + }); + + match &poll { + Poll::Ready(Some(Ok(batch))) => { + if batch.num_rows() > 0 { + break poll; + } else { + // Continue to poll input stream + } + } + Poll::Ready(Some(Err(_e))) => break poll, + Poll::Ready(None) => break poll, + Poll::Pending => break poll, + } + } + } + + /// Fetches from the batch + fn stream_limit(&mut self, batch: RecordBatch) -> Option<RecordBatch> { + // records time on drop + let _timer = self.baseline_metrics.elapsed_compute().timer(); + if self.fetch == 0 { + self.input = None; // Clear input so it can be dropped early + None + } else if batch.num_rows() < self.fetch { + // + self.fetch -= batch.num_rows(); + Some(batch) + } else if batch.num_rows() >= self.fetch { + let batch_rows = self.fetch; + self.fetch = 0; + self.input = None; // Clear input so it can be dropped early + + // It is guaranteed that batch_rows is <= batch.num_rows + Some(batch.slice(0, batch_rows)) + } else { + unreachable!() + } + } +} + +impl Stream for LimitStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let fetch_started = self.skip == 0; + let poll = match &mut self.input { + Some(input) => { + let poll = if fetch_started { + input.poll_next_unpin(cx) + } else { + self.poll_and_skip(cx) + }; + + poll.map(|x| match x { + Some(Ok(batch)) => Ok(self.stream_limit(batch)).transpose(), + other => other, + }) + } + // Input has been cleared + None => Poll::Ready(None), + }; + + self.baseline_metrics.record_poll(poll) + } +} + +impl RecordBatchStream for LimitStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::common::collect; + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::test; + + use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use arrow::array::RecordBatchOptions; + use arrow::compute::SortOptions; + use arrow::datatypes::Schema; + use datafusion_common::stats::Precision; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr}; + + #[tokio::test] + async fn limit() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + let num_partitions = 4; + let csv = test::scan_partitioned(num_partitions); + + // Input should have 4 partitions + assert_eq!(csv.output_partitioning().partition_count(), num_partitions); + + let limit = + GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), 0, Some(7)); + + // The result should contain 4 batches (one per input partition) + let iter = limit.execute(0, task_ctx)?; + let batches = collect(iter).await?; + + // There should be a total of 100 rows + let row_count: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(row_count, 7); + + Ok(()) + } + + #[tokio::test] + async fn limit_early_shutdown() -> Result<()> { + let batches = vec![ + test::make_partition(5), + test::make_partition(10), + test::make_partition(15), + test::make_partition(20), + test::make_partition(25), + ]; + let input = test::exec::TestStream::new(batches); + + let index = input.index(); + assert_eq!(index.value(), 0); + + // Limit of six needs to consume the entire first record batch + // (5 rows) and 1 row from the second (1 row) + let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let limit_stream = + LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics); + assert_eq!(index.value(), 0); + + let results = collect(Box::pin(limit_stream)).await.unwrap(); + let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); + // Only 6 rows should have been produced + assert_eq!(num_rows, 6); + + // Only the first two batches should be consumed + assert_eq!(index.value(), 2); + + Ok(()) + } + + #[tokio::test] + async fn limit_equals_batch_size() -> Result<()> { + let batches = vec![ + test::make_partition(6), + test::make_partition(6), + test::make_partition(6), + ]; + let input = test::exec::TestStream::new(batches); + + let index = input.index(); + assert_eq!(index.value(), 0); + + // Limit of six needs to consume the entire first record batch + // (6 rows) and stop immediately + let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let limit_stream = + LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics); + assert_eq!(index.value(), 0); + + let results = collect(Box::pin(limit_stream)).await.unwrap(); + let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); + // Only 6 rows should have been produced + assert_eq!(num_rows, 6); + + // Only the first batch should be consumed + assert_eq!(index.value(), 1); + + Ok(()) + } + + #[tokio::test] + async fn limit_no_column() -> Result<()> { + let batches = vec![ + make_batch_no_column(6), + make_batch_no_column(6), + make_batch_no_column(6), + ]; + let input = test::exec::TestStream::new(batches); + + let index = input.index(); + assert_eq!(index.value(), 0); + + // Limit of six needs to consume the entire first record batch + // (6 rows) and stop immediately + let baseline_metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let limit_stream = + LimitStream::new(Box::pin(input), 0, Some(6), baseline_metrics); + assert_eq!(index.value(), 0); + + let results = collect(Box::pin(limit_stream)).await.unwrap(); + let num_rows: usize = results.into_iter().map(|b| b.num_rows()).sum(); + // Only 6 rows should have been produced + assert_eq!(num_rows, 6); + + // Only the first batch should be consumed + assert_eq!(index.value(), 1); + + Ok(()) + } + + // Test cases for "skip" + async fn skip_and_fetch(skip: usize, fetch: Option<usize>) -> Result<usize> { + let task_ctx = Arc::new(TaskContext::default()); + + // 4 partitions @ 100 rows apiece + let num_partitions = 4; + let csv = test::scan_partitioned(num_partitions); + + assert_eq!(csv.output_partitioning().partition_count(), num_partitions); + + let offset = + GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch); + + // The result should contain 4 batches (one per input partition) + let iter = offset.execute(0, task_ctx)?; + let batches = collect(iter).await?; + Ok(batches.iter().map(|batch| batch.num_rows()).sum()) + } + + #[tokio::test] + async fn skip_none_fetch_none() -> Result<()> { + let row_count = skip_and_fetch(0, None).await?; + assert_eq!(row_count, 400); + Ok(()) + } + + #[tokio::test] + async fn skip_none_fetch_50() -> Result<()> { + let row_count = skip_and_fetch(0, Some(50)).await?; + assert_eq!(row_count, 50); + Ok(()) + } + + #[tokio::test] + async fn skip_3_fetch_none() -> Result<()> { + // There are total of 400 rows, we skipped 3 rows (offset = 3) + let row_count = skip_and_fetch(3, None).await?; + assert_eq!(row_count, 397); + Ok(()) + } + + #[tokio::test] + async fn skip_3_fetch_10_stats() -> Result<()> { + // There are total of 100 rows, we skipped 3 rows (offset = 3) + let row_count = skip_and_fetch(3, Some(10)).await?; + assert_eq!(row_count, 10); + Ok(()) + } + + #[tokio::test] + async fn skip_400_fetch_none() -> Result<()> { + let row_count = skip_and_fetch(400, None).await?; + assert_eq!(row_count, 0); + Ok(()) + } + + #[tokio::test] + async fn skip_400_fetch_1() -> Result<()> { + // There are a total of 400 rows + let row_count = skip_and_fetch(400, Some(1)).await?; + assert_eq!(row_count, 0); + Ok(()) + } + + #[tokio::test] + async fn skip_401_fetch_none() -> Result<()> { + // There are total of 400 rows, we skipped 401 rows (offset = 3) + let row_count = skip_and_fetch(401, None).await?; + assert_eq!(row_count, 0); + Ok(()) + } + + #[test] + fn replace_children_preserves_required_ordering() -> Result<()> { + let source = test::scan_partitioned(1); + let schema = source.schema(); + let ordering = LexOrdering::new(vec![PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }]); + + let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10)); + global.set_required_ordering(ordering.clone()); + let rebuilt = Arc::new(global).replace_children( + vec![test::scan_partitioned(1)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let rebuilt = rebuilt.downcast_ref::<GlobalLimitExec>().unwrap(); + assert_eq!(rebuilt.required_ordering(), &ordering); + + let mut local = LocalLimitExec::new(source, 10); + local.set_required_ordering(ordering.clone()); + let rebuilt = Arc::new(local).replace_children( + vec![test::scan_partitioned(1)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let rebuilt = rebuilt.downcast_ref::<LocalLimitExec>().unwrap(); + assert_eq!(rebuilt.required_ordering(), &ordering); + + Ok(()) + } + + #[test] + fn test_row_number_statistics_for_global_limit() -> Result<()> { + let row_count = row_number_statistics_for_global_limit(0, Some(10))?; + assert_eq!(row_count, Precision::Exact(10)); + + let row_count = row_number_statistics_for_global_limit(5, Some(10))?; + assert_eq!(row_count, Precision::Exact(10)); + + let row_count = row_number_statistics_for_global_limit(400, Some(10))?; + assert_eq!(row_count, Precision::Exact(0)); + + let row_count = row_number_statistics_for_global_limit(398, Some(10))?; + assert_eq!(row_count, Precision::Exact(2)); + + let row_count = row_number_statistics_for_global_limit(398, Some(1))?; + assert_eq!(row_count, Precision::Exact(1)); + + let row_count = row_number_statistics_for_global_limit(398, None)?; + assert_eq!(row_count, Precision::Exact(2)); + + let row_count = row_number_statistics_for_global_limit(0, Some(usize::MAX))?; + assert_eq!(row_count, Precision::Exact(400)); + + let row_count = row_number_statistics_for_global_limit(398, Some(usize::MAX))?; + assert_eq!(row_count, Precision::Exact(2)); + + let row_count = row_number_inexact_statistics_for_global_limit(0, Some(10))?; + assert_eq!(row_count, Precision::Inexact(10)); + + let row_count = row_number_inexact_statistics_for_global_limit(5, Some(10))?; + assert_eq!(row_count, Precision::Inexact(10)); + + // Input was Inexact, so an `nr <= skip` outcome must remain Inexact: + // the inexact estimate could be wrong, so we cannot promote 0 to + // Exact. + let row_count = row_number_inexact_statistics_for_global_limit(400, Some(10))?; + assert_eq!(row_count, Precision::Inexact(0)); + + let row_count = row_number_inexact_statistics_for_global_limit(398, Some(10))?; + assert_eq!(row_count, Precision::Inexact(2)); + + let row_count = row_number_inexact_statistics_for_global_limit(398, Some(1))?; + assert_eq!(row_count, Precision::Inexact(1)); + + let row_count = row_number_inexact_statistics_for_global_limit(398, None)?; + assert_eq!(row_count, Precision::Inexact(2)); + + let row_count = + row_number_inexact_statistics_for_global_limit(0, Some(usize::MAX))?; + assert_eq!(row_count, Precision::Inexact(400)); + + let row_count = + row_number_inexact_statistics_for_global_limit(398, Some(usize::MAX))?; + assert_eq!(row_count, Precision::Inexact(2)); + + Ok(()) + } + + #[test] + fn test_row_number_statistics_for_local_limit() -> Result<()> { + let row_count = row_number_statistics_for_local_limit(4, 10)?; + assert_eq!(row_count, Precision::Exact(10)); + + Ok(()) + } + + fn row_number_statistics_for_global_limit( + skip: usize, + fetch: Option<usize>, + ) -> Result<Precision<usize>> { + let num_partitions = 4; + let csv = test::scan_partitioned(num_partitions); + + assert_eq!(csv.output_partitioning().partition_count(), num_partitions); + + let offset = + GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch); + + Ok(StatisticsContext::new() + .compute(&offset, &StatisticsArgs::new())? + .num_rows) + } + + pub fn build_group_by( + input_schema: &SchemaRef, + columns: Vec<String>, + ) -> PhysicalGroupBy { + let mut group_by_expr: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![]; + for column in columns.iter() { + group_by_expr.push((col(column, input_schema).unwrap(), column.to_string())); + } + PhysicalGroupBy::new_single(group_by_expr.clone()) + } + + fn row_number_inexact_statistics_for_global_limit( + skip: usize, + fetch: Option<usize>, + ) -> Result<Precision<usize>> { + let num_partitions = 4; + let csv = test::scan_partitioned(num_partitions); + + assert_eq!(csv.output_partitioning().partition_count(), num_partitions); + + // Adding a "GROUP BY i" changes the input stats from Exact to Inexact. + let agg = AggregateExec::try_new( + AggregateMode::Final, + build_group_by(&csv.schema(), vec!["i".to_string()]), + vec![], + vec![], + Arc::clone(&csv), + Arc::clone(&csv.schema()), + )?; + let agg_exec: Arc<dyn ExecutionPlan> = Arc::new(agg); + + let offset = GlobalLimitExec::new( + Arc::new(CoalescePartitionsExec::new(agg_exec)), + skip, + fetch, + ); + + Ok(StatisticsContext::new() + .compute(&offset, &StatisticsArgs::new())? + .num_rows) + } + + fn row_number_statistics_for_local_limit( + num_partitions: usize, + fetch: usize, + ) -> Result<Precision<usize>> { + let csv = test::scan_partitioned(num_partitions); + + assert_eq!(csv.output_partitioning().partition_count(), num_partitions); + + let offset = LocalLimitExec::new(csv, fetch); + + Ok(StatisticsContext::new() + .compute(&offset, &StatisticsArgs::new())? + .num_rows) + } + + /// Return a RecordBatch with a single array with row_count sz + fn make_batch_no_column(sz: usize) -> RecordBatch { + let schema = Arc::new(Schema::empty()); + + let options = RecordBatchOptions::new().with_row_count(Option::from(sz)); + RecordBatch::try_new_with_options(schema, vec![], &options).unwrap() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/memory.rs b/native/vendor/datafusion-physical-plan/src/memory.rs new file mode 100644 index 0000000..0c77d7e --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/memory.rs
@@ -0,0 +1,993 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution plan for reading in-memory batches of data + +use std::any::Any; +use std::fmt; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use crate::coop::cooperative; +use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, +}; + +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_eq_or_internal_err, assert_or_internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; + +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use futures::Stream; +use parking_lot::RwLock; + +/// Iterator over batches +pub struct MemoryStream { + /// Vector of record batches + data: Vec<RecordBatch>, + /// Optional memory reservation bound to the data, freed on drop + reservation: Option<MemoryReservation>, + /// Schema representing the data + schema: SchemaRef, + /// Optional projection for which columns to load + projection: Option<Vec<usize>>, + /// Index into the data + index: usize, + /// The remaining number of rows to return. If None, all rows are returned + fetch: Option<usize>, +} + +impl MemoryStream { + /// Create an iterator for a vector of record batches + pub fn try_new( + data: Vec<RecordBatch>, + schema: SchemaRef, + projection: Option<Vec<usize>>, + ) -> Result<Self> { + Ok(Self { + data, + reservation: None, + schema, + projection, + index: 0, + fetch: None, + }) + } + + /// Set the memory reservation for the data + pub fn with_reservation(mut self, reservation: MemoryReservation) -> Self { + self.reservation = Some(reservation); + self + } + + /// Set the number of rows to produce + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } +} + +impl Stream for MemoryStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.index >= self.data.len() { + return Poll::Ready(None); + } + self.index += 1; + let batch = &self.data[self.index - 1]; + // return just the columns requested + let batch = match self.projection.as_ref() { + Some(columns) => batch.project(columns)?, + None => batch.clone(), + }; + + // MemoryStream advertises `self.schema`, therefore emitted RecordBatches + // must conform to it when batches were provided with stricter nested types + // (e.g. MemTable accepts stricter batches via Schema::contains). + let batch = if batch.schema().as_ref() != self.schema.as_ref() + && self.schema.contains(batch.schema().as_ref()) + { + datafusion_common::nested_struct::adapt_batch_to_schema(batch, &self.schema)? + } else { + batch + }; + + let Some(&fetch) = self.fetch.as_ref() else { + return Poll::Ready(Some(Ok(batch))); + }; + if fetch == 0 { + return Poll::Ready(None); + } + + let batch = if batch.num_rows() > fetch { + batch.slice(0, fetch) + } else { + batch + }; + self.fetch = Some(fetch - batch.num_rows()); + Poll::Ready(Some(Ok(batch))) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (self.data.len(), Some(self.data.len())) + } +} + +impl RecordBatchStream for MemoryStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +pub trait LazyBatchGenerator: Send + Sync + fmt::Debug + fmt::Display { + /// Returns the generator as [`Any`] so that it can be + /// downcast to a specific implementation. + fn as_any(&self) -> &dyn Any; + + fn boundedness(&self) -> Boundedness { + Boundedness::Bounded + } + + /// Generate the next batch, return `None` when no more batches are available + fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>>; + + /// Returns a new instance with the state reset. + fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>>; +} + +/// Execution plan for lazy in-memory batches of data +/// +/// This plan generates output batches lazily, it doesn't have to buffer all batches +/// in memory up front (compared to `MemorySourceConfig`), thus consuming constant memory. +pub struct LazyMemoryExec { + /// Schema representing the data + schema: SchemaRef, + /// Optional projection for which columns to load + projection: Option<Vec<usize>>, + /// Functions to generate batches for each partition + batch_generators: Vec<Arc<RwLock<dyn LazyBatchGenerator>>>, + /// Plan properties cache storing equivalence properties, partitioning, and execution mode + cache: Arc<PlanProperties>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, +} + +impl LazyMemoryExec { + /// Create a new lazy memory execution plan + pub fn try_new( + schema: SchemaRef, + generators: Vec<Arc<RwLock<dyn LazyBatchGenerator>>>, + ) -> Result<Self> { + let boundedness = generators + .iter() + .map(|g| g.read().boundedness()) + .reduce(|acc, b| match acc { + Boundedness::Bounded => b, + Boundedness::Unbounded { + requires_infinite_memory, + } => { + let acc_infinite_memory = requires_infinite_memory; + match b { + Boundedness::Bounded => acc, + Boundedness::Unbounded { + requires_infinite_memory, + } => Boundedness::Unbounded { + requires_infinite_memory: requires_infinite_memory + || acc_infinite_memory, + }, + } + } + }) + .unwrap_or(Boundedness::Bounded); + + let cache = PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::RoundRobinBatch(generators.len()), + EmissionType::Incremental, + boundedness, + ) + .with_scheduling_type(SchedulingType::Cooperative) + .into(); + + Ok(Self { + schema, + projection: None, + batch_generators: generators, + cache, + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + pub fn with_projection(mut self, projection: Option<Vec<usize>>) -> Self { + match projection.as_ref() { + Some(columns) => { + let projected = Arc::new(self.schema.project(columns).unwrap()); + Arc::make_mut(&mut self.cache).set_eq_properties( + EquivalenceProperties::new(Arc::clone(&projected)), + ); + self.schema = projected; + self.projection = projection; + self + } + _ => self, + } + } + + pub fn try_set_partitioning(&mut self, partitioning: Partitioning) -> Result<()> { + let partition_count = partitioning.partition_count(); + let generator_count = self.batch_generators.len(); + assert_eq_or_internal_err!( + partition_count, + generator_count, + "Partition count must match generator count: {} != {}", + partition_count, + generator_count + ); + Arc::make_mut(&mut self.cache).partitioning = partitioning; + Ok(()) + } + + pub fn add_ordering(&mut self, ordering: impl IntoIterator<Item = PhysicalSortExpr>) { + Arc::make_mut(&mut self.cache) + .eq_properties + .add_orderings(std::iter::once(ordering)); + } + + /// Get the batch generators + pub fn generators(&self) -> &Vec<Arc<RwLock<dyn LazyBatchGenerator>>> { + &self.batch_generators + } +} + +impl fmt::Debug for LazyMemoryExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("LazyMemoryExec") + .field("schema", &self.schema) + .field("batch_generators", &self.batch_generators) + .finish() + } +} + +impl DisplayAs for LazyMemoryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "LazyMemoryExec: partitions={}, batch_generators=[{}]", + self.batch_generators.len(), + self.batch_generators + .iter() + .map(|g| g.read().to_string()) + .collect::<Vec<_>>() + .join(", ") + ) + } + DisplayFormatType::TreeRender => { + //TODO: remove batch_size, add one line per generator + writeln!( + f, + "batch_generators={}", + self.batch_generators + .iter() + .map(|g| g.read().to_string()) + .collect::<Vec<String>>() + .join(", ") + )?; + Ok(()) + } + } + } +} + +impl ExecutionPlan for LazyMemoryExec { + fn name(&self) -> &'static str { + "LazyMemoryExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + assert_or_internal_err!( + children.is_empty(), + "Children cannot be replaced in LazyMemoryExec" + ); + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + assert_or_internal_err!( + partition < self.batch_generators.len(), + "Invalid partition {} for LazyMemoryExec with {} partitions", + partition, + self.batch_generators.len() + ); + + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + + // Create a fresh generator via reset_state() so that each execute() + // call produces an independent stream starting from the beginning. + let generator = self.batch_generators[partition].read().reset_state(); + + let stream = LazyMemoryStream { + schema: Arc::clone(&self.schema), + projection: self.projection.clone(), + generator, + baseline_metrics, + }; + Ok(Box::pin(cooperative(stream))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + let generators = self + .generators() + .iter() + .map(|g| g.read().reset_state()) + .collect::<Vec<_>>(); + Ok(Arc::new(LazyMemoryExec { + schema: Arc::clone(&self.schema), + batch_generators: generators, + cache: Arc::clone(&self.cache), + metrics: ExecutionPlanMetricsSet::new(), + projection: self.projection.clone(), + })) + } +} + +/// Stream that generates record batches on demand +pub struct LazyMemoryStream { + schema: SchemaRef, + /// Optional projection for which columns to load + projection: Option<Vec<usize>>, + /// Generator to produce batches + /// + /// Note: Idiomatically, DataFusion uses plan-time parallelism - each stream + /// should have a unique `LazyBatchGenerator`. Use RepartitionExec or + /// construct multiple `LazyMemoryStream`s during planning to enable + /// parallel execution. + /// Sharing generators between streams should be used with caution. + generator: Arc<RwLock<dyn LazyBatchGenerator>>, + /// Execution metrics + baseline_metrics: BaselineMetrics, +} + +impl Stream for LazyMemoryStream { + type Item = Result<RecordBatch>; + + fn poll_next( + self: std::pin::Pin<&mut Self>, + _: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let _timer_guard = self.baseline_metrics.elapsed_compute().timer(); + let batch = self.generator.write().generate_next_batch(); + + let poll = match batch { + Ok(Some(batch)) => { + // return just the columns requested + let batch = match self.projection.as_ref() { + Some(columns) => batch.project(columns)?, + None => batch, + }; + Poll::Ready(Some(Ok(batch))) + } + Ok(None) => Poll::Ready(None), + Err(e) => Poll::Ready(Some(Err(e))), + }; + + self.baseline_metrics.record_poll(poll) + } +} + +impl RecordBatchStream for LazyMemoryStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod lazy_memory_tests { + use super::*; + use crate::common::collect; + use arrow::array::Int64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use futures::StreamExt; + + #[derive(Debug, Clone)] + struct TestGenerator { + counter: i64, + max_batches: i64, + batch_size: usize, + schema: SchemaRef, + } + + impl fmt::Display for TestGenerator { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "TestGenerator: counter={}, max_batches={}, batch_size={}", + self.counter, self.max_batches, self.batch_size + ) + } + } + + impl LazyBatchGenerator for TestGenerator { + fn as_any(&self) -> &dyn Any { + self + } + + fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>> { + if self.counter >= self.max_batches { + return Ok(None); + } + + let array = Int64Array::from_iter_values( + (self.counter * self.batch_size as i64) + ..(self.counter * self.batch_size as i64 + self.batch_size as i64), + ); + self.counter += 1; + Ok(Some(RecordBatch::try_new( + Arc::clone(&self.schema), + vec![Arc::new(array)], + )?)) + } + + fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> { + Arc::new(RwLock::new(TestGenerator { + counter: 0, + max_batches: self.max_batches, + batch_size: self.batch_size, + schema: Arc::clone(&self.schema), + })) + } + } + + #[tokio::test] + async fn test_lazy_memory_exec() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let generator = TestGenerator { + counter: 0, + max_batches: 3, + batch_size: 2, + schema: Arc::clone(&schema), + }; + + let exec = + LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?; + + // Test schema + assert_eq!(exec.schema().fields().len(), 1); + assert_eq!(exec.schema().field(0).name(), "a"); + + // Test execution + let stream = exec.execute(0, Arc::new(TaskContext::default()))?; + let batches: Vec<_> = stream.collect::<Vec<_>>().await; + + assert_eq!(batches.len(), 3); + + // Verify batch contents + let batch0 = batches[0].as_ref().unwrap(); + let array0 = batch0 + .column(0) + .as_any() + .downcast_ref::<Int64Array>() + .unwrap(); + assert_eq!(array0.values(), &[0, 1]); + + let batch1 = batches[1].as_ref().unwrap(); + let array1 = batch1 + .column(0) + .as_any() + .downcast_ref::<Int64Array>() + .unwrap(); + assert_eq!(array1.values(), &[2, 3]); + + let batch2 = batches[2].as_ref().unwrap(); + let array2 = batch2 + .column(0) + .as_any() + .downcast_ref::<Int64Array>() + .unwrap(); + assert_eq!(array2.values(), &[4, 5]); + + Ok(()) + } + + /// Verify that calling execute(0) twice on the same LazyMemoryExec + /// produces independent streams with the same data. + #[tokio::test] + async fn test_lazy_memory_exec_multiple_executions_are_independent() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let generator = TestGenerator { + counter: 0, + max_batches: 3, + batch_size: 2, + schema: Arc::clone(&schema), + }; + + let exec = + LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?; + let task_ctx = Arc::new(TaskContext::default()); + + // First execution — consume all batches + let batches_1 = collect(exec.execute(0, Arc::clone(&task_ctx))?).await?; + let total_rows_1: usize = batches_1.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows_1, 6); + + // Second execution — should produce the same data, not continue + // from where the first execution left off + let batches_2 = collect(exec.execute(0, Arc::clone(&task_ctx))?).await?; + let total_rows_2: usize = batches_2.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows_2, 6); + + // Verify contents are identical + for (b1, b2) in batches_1.iter().zip(batches_2.iter()) { + assert_eq!(b1, b2); + } + + Ok(()) + } + + #[tokio::test] + async fn test_lazy_memory_exec_invalid_partition() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let generator = TestGenerator { + counter: 0, + max_batches: 1, + batch_size: 1, + schema: Arc::clone(&schema), + }; + + let exec = + LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?; + + // Test invalid partition + let result = exec.execute(1, Arc::new(TaskContext::default())); + + // partition is 0-indexed, so there only should be partition 0 + assert!(matches!( + result, + Err(e) if e.to_string().contains("Invalid partition 1 for LazyMemoryExec with 1 partitions") + )); + + Ok(()) + } + + #[tokio::test] + async fn test_generate_series_metrics_integration() -> Result<()> { + // Test LazyMemoryExec metrics with different configurations + let test_cases = vec![ + (10, 2, 10), // 10 rows, batch size 2, expected 10 rows + (100, 10, 100), // 100 rows, batch size 10, expected 100 rows + (5, 1, 5), // 5 rows, batch size 1, expected 5 rows + ]; + + for (total_rows, batch_size, expected_rows) in test_cases { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let generator = TestGenerator { + counter: 0, + max_batches: (total_rows + batch_size - 1) / batch_size, // ceiling division + batch_size: batch_size as usize, + schema: Arc::clone(&schema), + }; + + let exec = + LazyMemoryExec::try_new(schema, vec![Arc::new(RwLock::new(generator))])?; + let task_ctx = Arc::new(TaskContext::default()); + + let stream = exec.execute(0, task_ctx)?; + let batches = collect(stream).await?; + + // Verify metrics exist with actual expected numbers + let metrics = exec.metrics().unwrap(); + + // Count actual rows returned + let actual_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(actual_rows, expected_rows); + + // Verify metrics match actual output + assert_eq!(metrics.output_rows().unwrap(), expected_rows); + assert!(metrics.elapsed_compute().unwrap() > 0); + } + + Ok(()) + } + + #[tokio::test] + async fn test_lazy_memory_exec_reset_state() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let generator = TestGenerator { + counter: 0, + max_batches: 3, + batch_size: 2, + schema: Arc::clone(&schema), + }; + + let exec = Arc::new(LazyMemoryExec::try_new( + schema, + vec![Arc::new(RwLock::new(generator))], + )?); + let stream = exec.execute(0, Arc::new(TaskContext::default()))?; + let batches = collect(stream).await?; + + let exec_reset = exec.reset_state()?; + let stream = exec_reset.execute(0, Arc::new(TaskContext::default()))?; + let batches_reset = collect(stream).await?; + + // if the reset_state is not correct, the batches_reset will be empty + assert_eq!(batches, batches_reset); + + Ok(()) + } + + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema() -> Result<()> { + use arrow::array::{ArrayRef, BooleanArray, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + // Declared schema expects nullable struct field colA + let declared_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, true)]); + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::Struct(declared_fields), + false, + )])); + + // Runtime batch has stricter non-nullable struct field colA + let source_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let source_schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::Struct(source_fields.clone()), + false, + )])); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + source_fields, + vec![Arc::new(BooleanArray::from(vec![true, false]))], + None, + )); + let stricter_batch = RecordBatch::try_new(source_schema, vec![struct_array])?; + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), declared_schema); + + let struct_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::<StructArray>() + .unwrap(); + assert!(struct_col.fields()[0].is_nullable()); + let bool_child = struct_col + .column(0) + .as_any() + .downcast_ref::<BooleanArray>() + .unwrap(); + assert!(bool_child.value(0)); + assert!(!bool_child.value(1)); + + Ok(()) + } + + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_with_projection() + -> Result<()> { + use arrow::array::{ArrayRef, BooleanArray, Int32Array, StructArray}; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + // Declared full schema: col a (Int32), col b (Struct<colA: nullable Boolean>) + let declared_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, true)]); + let full_declared_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Struct(declared_fields), false), + ])); + + // Projected schema for column "b" (projection = [1]) + let projected_schema = Arc::new(full_declared_schema.project(&[1])?); + + // Runtime batch has stricter struct + let source_fields = + Fields::from(vec![Field::new("colA", DataType::Boolean, false)]); + let source_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Struct(source_fields.clone()), false), + ])); + + let struct_array: ArrayRef = Arc::new(StructArray::new( + source_fields, + vec![Arc::new(BooleanArray::from(vec![true, false]))], + None, + )); + let stricter_batch = RecordBatch::try_new( + source_schema, + vec![Arc::new(Int32Array::from(vec![10, 20])), struct_array], + )?; + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&projected_schema), + Some(vec![1]), + )?; + + assert_eq!(stream.schema(), projected_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), projected_schema); + assert_eq!(emitted_batch.num_columns(), 1); + + let struct_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::<StructArray>() + .unwrap(); + assert!(struct_col.fields()[0].is_nullable()); + let bool_child = struct_col + .column(0) + .as_any() + .downcast_ref::<BooleanArray>() + .unwrap(); + assert!(bool_child.value(0)); + assert!(!bool_child.value(1)); + + Ok(()) + } + + /// Regression for the Union reconstruction path at the `MemoryStream` + /// producer boundary: a declared nullable Union child vs a stricter + /// non-nullable runtime child. + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_union() -> Result<()> + { + use arrow::array::{Array, ArrayRef, Float64Array, Int32Array, UnionArray}; + use arrow::buffer::ScalarBuffer; + use arrow::datatypes::{DataType, Field, Schema, UnionFields, UnionMode}; + use futures::StreamExt; + + let declared_union_fields = UnionFields::try_new( + vec![0_i8, 1], + vec![ + Field::new("i", DataType::Int32, true), + Field::new("f", DataType::Float64, true), + ], + )?; + let declared_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(declared_union_fields, UnionMode::Dense), + false, + )])); + + let source_union_fields = UnionFields::try_new( + vec![0_i8, 1], + vec![ + Field::new("i", DataType::Int32, false), + Field::new("f", DataType::Float64, false), + ], + )?; + let source_schema = Arc::new(Schema::new(vec![Field::new( + "u", + DataType::Union(source_union_fields.clone(), UnionMode::Dense), + false, + )])); + + let type_ids = ScalarBuffer::from(vec![0_i8, 1, 0]); + let offsets = ScalarBuffer::from(vec![0_i32, 0, 1]); + let union_array: ArrayRef = Arc::new(UnionArray::try_new( + source_union_fields, + type_ids, + Some(offsets), + vec![ + Arc::new(Int32Array::from(vec![10, 20])), + Arc::new(Float64Array::from(vec![1.5])), + ], + )?); + let stricter_batch = RecordBatch::try_new(source_schema, vec![union_array])?; + + assert!(declared_schema.contains(stricter_batch.schema().as_ref())); + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), stream.schema()); + assert_eq!(emitted_batch.schema(), declared_schema); + + let union_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::<UnionArray>() + .unwrap(); + assert_eq!(union_col.len(), 3); + assert_eq!(union_col.type_id(0), 0); + assert_eq!(union_col.type_id(1), 1); + assert_eq!(union_col.type_id(2), 0); + let i_child = union_col + .child(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(i_child.values(), &[10, 20]); + + Ok(()) + } + + /// Regression for a contained `Map<.., Struct>` whose runtime nested field + /// is non-nullable while the declared nested field is nullable. + #[tokio::test] + async fn test_memory_stream_emitted_batch_matches_declared_schema_map_of_struct() + -> Result<()> { + use arrow::array::{ + Array, ArrayRef, Int32Array, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Fields, Schema}; + use futures::StreamExt; + + fn map_field(value_child_nullable: bool) -> Field { + let value_struct = DataType::Struct(Fields::from(vec![Field::new( + "v", + DataType::Int32, + value_child_nullable, + )])); + let entries = Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", value_struct, true), + ])), + false, + ); + Field::new("m", DataType::Map(Arc::new(entries), false), true) + } + + let declared_schema = Arc::new(Schema::new(vec![map_field(true)])); + let source_schema = Arc::new(Schema::new(vec![map_field(false)])); + + let value_fields = Fields::from(vec![Field::new("v", DataType::Int32, false)]); + let values_struct = StructArray::new( + value_fields, + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + None, + ); + let entries = StructArray::new( + Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", values_struct.data_type().clone(), true), + ]), + vec![ + Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef, + Arc::new(values_struct) as ArrayRef, + ], + None, + ); + let DataType::Map(source_entries_field, _) = source_schema.field(0).data_type() + else { + unreachable!("map field") + }; + let map_array: ArrayRef = Arc::new(MapArray::try_new( + Arc::clone(source_entries_field), + OffsetBuffer::new(vec![0, 2, 3].into()), + entries, + None, + false, + )?); + let stricter_batch = RecordBatch::try_new(source_schema, vec![map_array])?; + + // The stricter batch is accepted by `MemTable::try_new`-style checks. + assert!(declared_schema.contains(stricter_batch.schema().as_ref())); + + let mut stream = MemoryStream::try_new( + vec![stricter_batch], + Arc::clone(&declared_schema), + None, + )?; + + assert_eq!(stream.schema(), declared_schema); + + let emitted_batch = stream.next().await.unwrap()?; + assert_eq!(emitted_batch.schema(), stream.schema()); + assert_eq!(emitted_batch.schema(), declared_schema); + + let map_col = emitted_batch + .column(0) + .as_any() + .downcast_ref::<MapArray>() + .unwrap(); + assert_eq!(map_col.len(), 2); + let values = map_col + .values() + .as_any() + .downcast_ref::<StructArray>() + .unwrap(); + assert!(values.fields()[0].is_nullable()); + let ints = values + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(ints.values(), &[1, 2, 3]); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/metrics.rs b/native/vendor/datafusion-physical-plan/src/metrics.rs new file mode 100644 index 0000000..fe17cbd --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/metrics.rs
@@ -0,0 +1,21 @@ +// 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. + +//! Metrics live in `datafusion-physical-expr-common`; this module re-exports +//! them to keep the public APIs stable. + +pub use datafusion_physical_expr_common::metrics::*;
diff --git a/native/vendor/datafusion-physical-plan/src/operator_statistics/mod.rs b/native/vendor/datafusion-physical-plan/src/operator_statistics/mod.rs new file mode 100644 index 0000000..16b89e9 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/operator_statistics/mod.rs
@@ -0,0 +1,2342 @@ +// 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. + +//! Pluggable statistics propagation for physical plans. +//! +//! This module provides an extensible mechanism for computing statistics +//! on [`ExecutionPlan`] nodes, following the chain of responsibility pattern +//! similar to `RelationPlanner` for SQL parsing. +//! +//! # Overview +//! +//! The default implementation delegates to each operator's built-in +//! `partition_statistics`. Users can register custom [`StatisticsProvider`] +//! implementations to: +//! +//! 1. Provide statistics for custom [`ExecutionPlan`] implementations +//! 2. Override default estimation with advanced approaches (e.g., histograms) +//! 3. Plug in domain-specific knowledge for better cardinality estimation +//! +//! # Architecture +//! +//! - [`StatisticsProvider`]: Chain element that computes statistics for specific operators +//! - [`StatisticsRegistry`]: Chains providers, lives in SessionState +//! - [`ExtendedStatistics`]: Statistics with type-safe custom extensions +//! +//! # Built-in Providers +//! +//! The following providers are included and can be registered in this order: +//! +//! 1. [`FilterStatisticsProvider`] - selectivity-based filter estimation +//! 2. [`ProjectionStatisticsProvider`] - column mapping through projections +//! 3. [`PassthroughStatisticsProvider`] - passthrough for cardinality-preserving operators +//! 4. [`AggregateStatisticsProvider`] - NDV-based GROUP BY cardinality estimation +//! 5. [`JoinStatisticsProvider`] - NDV-based join output estimation (hash, sort-merge, cross) +//! 6. [`LimitStatisticsProvider`] - caps output at the fetch limit (local and global) +//! 7. [`UnionStatisticsProvider`] - sums input row counts +//! 8. [`DefaultStatisticsProvider`] - fallback to `partition_statistics(None)` +//! +//! # Relationship to [#20184](https://github.com/apache/datafusion/issues/20184) +//! +//! This module performs its own bottom-up tree walk in [`StatisticsRegistry::compute`], +//! separate from the walk optimizer rules do via `transform_up`. This means existing +//! rules that call `partition_statistics` directly bypass the registry. +//! +//! [#20184](https://github.com/apache/datafusion/issues/20184) adds a `child_stats` +//! parameter to `partition_statistics`. Once it lands, the registry can feed enriched +//! **base** [`Statistics`] into operators' built-in `partition_statistics` calls, +//! removing redundancy for the base-stats path (row counts, column stats). However, +//! the separate registry walk is still required for [`ExtendedStatistics`] extension +//! propagation: `partition_statistics` returns `Arc<Statistics>`, so extensions +//! (histograms, sketches, etc.) are stripped at that boundary and can only flow +//! through the registry walk. +//! +//! If [`Statistics`] itself were extended to carry a type-erased extension map +//! (similar to [`ExtendedStatistics`]), the registry walk could be dropped entirely: +//! extensions would flow naturally through `partition_statistics(child_stats)` and +//! the registry would become a pure chain-of-responsibility on top of the existing +//! traversal with no separate walk needed. +//! +//! # Example +//! +//! ```ignore +//! use datafusion_physical_plan::operator_statistics::*; +//! +//! // Create registry with default provider +//! let mut registry = StatisticsRegistry::new(); +//! +//! // Register custom provider (higher priority) +//! registry.register(Arc::new(MyHistogramProvider)); +//! +//! // Compute statistics through the chain +//! let stats = registry.compute(plan.as_ref())?; +//! ``` + +use std::fmt::{self, Debug}; +use std::sync::Arc; + +use datafusion_common::extensions::Extensions; +use datafusion_common::stats::Precision; +use datafusion_common::{Result, Statistics}; + +use crate::ExecutionPlan; +use crate::statistics::{StatisticsArgs, StatisticsContext}; + +// ============================================================================ +// ExtendedStatistics: Statistics with type-safe extensions +// ============================================================================ + +/// Statistics with support for custom extensions. +/// +/// Wraps the standard [`Statistics`] and adds a type-erased extension map +/// for custom statistics like histograms, sketches, or domain-specific metadata. +/// +/// # Example +/// +/// ```ignore +/// // Define a custom statistics extension +/// #[derive(Debug, Clone)] +/// struct HistogramStats { +/// buckets: Vec<(i64, i64, usize)>, // (min, max, count) +/// } +/// +/// // Set extension in a planner +/// let mut stats = ExtendedStatistics::from(base_stats); +/// stats.set_extension(HistogramStats { buckets: vec![] }); +/// +/// // Retrieve in a consumer +/// if let Some(hist) = stats.get_extension::<HistogramStats>() { +/// // Use histogram for better estimation +/// } +/// ``` +#[derive(Debug, Clone, Default)] +pub struct ExtendedStatistics { + /// Standard statistics (num_rows, byte_size, column stats) + base: Arc<Statistics>, + /// Type-erased extensions for custom statistics + extensions: Extensions, +} + +impl ExtendedStatistics { + /// Create new ExtendedStatistics wrapping owned statistics. + pub fn new(base: Statistics) -> Self { + Self { + base: Arc::new(base), + extensions: Extensions::new(), + } + } + + /// Create new ExtendedStatistics from an [`Arc<Statistics>`]. + pub fn new_arc(base: Arc<Statistics>) -> Self { + Self { + base, + extensions: Extensions::new(), + } + } + + /// Returns a reference to the base [`Statistics`]. + pub fn base(&self) -> &Statistics { + &self.base + } + + /// Returns a reference to the underlying [`Arc<Statistics>`]. + pub fn base_arc(&self) -> &Arc<Statistics> { + &self.base + } + + /// Get a reference to a custom statistics extension by type. + pub fn get_extension<T: 'static + Send + Sync>(&self) -> Option<&T> { + self.extensions.get::<T>() + } + + /// Set a custom statistics extension. + pub fn set_extension<T: 'static + Send + Sync>(&mut self, value: T) { + self.extensions.insert(value); + } + + /// Check if an extension of the given type exists. + pub fn has_extension<T: 'static + Send + Sync>(&self) -> bool { + self.extensions.contains::<T>() + } + + /// Merge extensions from another ExtendedStatistics (other's extensions take precedence). + pub fn merge_extensions(&mut self, other: &ExtendedStatistics) { + self.extensions.merge(&other.extensions); + } +} + +impl From<Statistics> for ExtendedStatistics { + fn from(base: Statistics) -> Self { + Self::new(base) + } +} + +impl From<Arc<Statistics>> for ExtendedStatistics { + fn from(base: Arc<Statistics>) -> Self { + Self::new_arc(base) + } +} + +impl From<ExtendedStatistics> for Statistics { + fn from(extended: ExtendedStatistics) -> Self { + Arc::unwrap_or_clone(extended.base) + } +} + +// ============================================================================ +// StatisticsProvider trait and registry +// ============================================================================ + +/// Result of attempting to compute statistics with a [`StatisticsProvider`]. +#[derive(Debug)] +pub enum StatisticsResult { + /// Statistics were computed by this provider + Computed(ExtendedStatistics), + /// This provider doesn't handle this operator; delegate to next in chain + Delegate, +} + +/// Customize statistics computation for [`ExecutionPlan`] nodes. +/// +/// Implementations can handle specific operator types or override default +/// estimation logic. The chain of providers is traversed until one returns +/// [`StatisticsResult::Computed`]. +/// +/// # Implementing a Custom Provider +/// +/// ```ignore +/// #[derive(Debug)] +/// struct MyStatisticsProvider; +/// +/// impl StatisticsProvider for MyStatisticsProvider { +/// fn compute_statistics( +/// &self, +/// plan: &dyn ExecutionPlan, +/// child_stats: &[ExtendedStatistics], +/// ) -> Result<StatisticsResult> { +/// if let Some(my_exec) = plan.downcast_ref::<MyCustomExec>() { +/// // Custom logic for MyCustomExec +/// Ok(StatisticsResult::Computed(/* ... */)) +/// } else { +/// // Let next provider handle it +/// Ok(StatisticsResult::Delegate) +/// } +/// } +/// } +/// ``` +pub trait StatisticsProvider: Debug + Send + Sync { + /// Compute statistics for an [`ExecutionPlan`] node. + /// + /// # Arguments + /// * `plan` - The execution plan node to compute statistics for + /// * `child_stats` - Extended statistics already computed for child nodes, + /// in the same order as `plan.children()`. Empty for leaf nodes. + /// + /// # Returns + /// * `StatisticsResult::Computed(stats)` - Short-circuits the chain + /// * `StatisticsResult::Delegate` - Passes to next provider in chain + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult>; +} + +/// Default statistics provider that delegates to each operator's built-in +/// `partition_statistics` implementation. +#[derive(Debug, Default)] +pub struct DefaultStatisticsProvider; + +impl StatisticsProvider for DefaultStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + _child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; + Ok(StatisticsResult::Computed(ExtendedStatistics::new_arc( + base, + ))) + } +} + +/// Registry that chains [`StatisticsProvider`] implementations. +/// +/// The registry is a stateless provider chain: it holds no mutable state +/// and is cheaply `Clone`able / `Send` / `Sync`. +#[derive(Clone)] +pub struct StatisticsRegistry { + providers: Vec<Arc<dyn StatisticsProvider>>, +} + +impl Debug for StatisticsRegistry { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "StatisticsRegistry({} providers)", self.providers.len()) + } +} + +impl Default for StatisticsRegistry { + fn default() -> Self { + Self::new() + } +} + +impl StatisticsRegistry { + /// Create a new empty registry. + /// + /// With no providers, `compute()` falls back to each plan node's + /// built-in `partition_statistics()`. Register providers to enhance + /// statistics (e.g., inject NDV, use histograms). + pub fn new() -> Self { + Self { + providers: Vec::new(), + } + } + + /// Create a registry with the given provider chain. + pub fn with_providers(providers: Vec<Arc<dyn StatisticsProvider>>) -> Self { + Self { providers } + } + + /// Create a registry pre-loaded with the standard built-in providers. + /// + /// Provider order (first match wins): + /// 1. [`FilterStatisticsProvider`] + /// 2. [`ProjectionStatisticsProvider`] + /// 3. [`PassthroughStatisticsProvider`] + /// 4. [`AggregateStatisticsProvider`] + /// 5. [`JoinStatisticsProvider`] + /// 6. [`LimitStatisticsProvider`] + /// 7. [`UnionStatisticsProvider`] + /// 8. [`DefaultStatisticsProvider`] + pub fn default_with_builtin_providers() -> Self { + Self::with_providers(vec![ + Arc::new(FilterStatisticsProvider), + Arc::new(ProjectionStatisticsProvider), + Arc::new(PassthroughStatisticsProvider), + Arc::new(AggregateStatisticsProvider), + Arc::new(JoinStatisticsProvider), + Arc::new(LimitStatisticsProvider), + Arc::new(UnionStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]) + } + + /// Register a provider at the front of the chain (higher priority). + pub fn register(&mut self, provider: Arc<dyn StatisticsProvider>) { + self.providers.insert(0, provider); + } + + /// Returns the current provider chain. + pub fn providers(&self) -> &[Arc<dyn StatisticsProvider>] { + &self.providers + } + + /// Compute extended statistics for a plan through the provider chain. + /// + /// Performs a bottom-up tree walk: child statistics are computed recursively + /// and passed to providers, mirroring how `partition_statistics` composes + /// operators. Once [#20184](https://github.com/apache/datafusion/issues/20184) + /// lands, the registry can feed enriched base stats directly into + /// `partition_statistics(child_stats)`, removing the need for a separate walk. + /// + /// If no providers are registered, falls back to the plan's built-in + /// `partition_statistics(None)` with no overhead. + pub fn compute(&self, plan: &dyn ExecutionPlan) -> Result<ExtendedStatistics> { + // Fast path: no providers registered, skip the walk entirely + if self.providers.is_empty() { + let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; + return Ok(ExtendedStatistics::new_arc(base)); + } + + let children = plan.children(); + + // For leaf nodes, try providers with empty child stats. + // For non-leaf nodes, recursively compute enhanced child stats first. + let child_stats: Vec<ExtendedStatistics> = if children.is_empty() { + Vec::new() + } else { + children + .iter() + .map(|child| self.compute(child.as_ref())) + .collect::<Result<Vec<_>>>()? + }; + + for provider in &self.providers { + match provider.compute_statistics(plan, &child_stats)? { + StatisticsResult::Computed(stats) => return Ok(stats), + StatisticsResult::Delegate => continue, + } + } + // Fallback: use plan's built-in stats + let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; + Ok(ExtendedStatistics::new_arc(base)) + } + + /// Compute statistics and return only the base Statistics (no extensions). + /// + /// Convenience method for callers that don't need extensions. + pub fn compute_base(&self, plan: &dyn ExecutionPlan) -> Result<Statistics> { + Ok(self.compute(plan)?.base().clone()) + } +} + +// ============================================================================ +// Statistics Utility Functions +// ============================================================================ + +/// Estimate the number of distinct values when sampling from a population. +/// +/// Given a domain with `domain_size` distinct values and `num_selected` rows +/// sampled/filtered from it, estimates how many distinct values will appear +/// in the sample. +/// +/// Uses the formula: `Expected distinct = N * [1 - (1 - 1/N)^n]` +/// +/// # References +/// +/// Based on Calcite's `RelMdUtil.numDistinctVals()`: +/// <https://github.com/apache/calcite/blob/main/core/src/main/java/org/apache/calcite/rel/metadata/RelMdUtil.java> +pub fn num_distinct_vals(domain_size: usize, num_selected: usize) -> usize { + if domain_size == 0 || num_selected == 0 { + return 0; + } + + if num_selected >= domain_size { + return domain_size; + } + + let n = domain_size as f64; + let k = num_selected as f64; + + // For large n, (1-1/n).powf(k) loses precision because the base is near + // 1.0; use the equivalent exp(-k/n) form which is numerically stable. + // Threshold matches Calcite's RelMdUtil.numDistinctVals(). + let expected = if domain_size > 1000 { + n * (1.0 - (-k / n).exp()) + } else { + n * (1.0 - (1.0 - 1.0 / n).powf(k)) + }; + + let result = expected.round() as usize; + result.clamp(1, domain_size) +} + +/// Estimate NDV after applying a selectivity factor (filtering). +/// +/// When filtering rows, each distinct value has multiple rows. If a value +/// appears `k` times, the probability it survives the filter is `1 - (1-s)^k` +/// where `s` is the selectivity. +/// +/// Assuming uniform distribution (each value appears `rows/ndv` times): +/// ```text +/// NDV_after ~ NDV_before * [1 - (1 - selectivity)^(rows/NDV)] +/// ``` +pub fn ndv_after_selectivity( + original_ndv: usize, + original_rows: usize, + selectivity: f64, +) -> usize { + if selectivity <= 0.0 || original_ndv == 0 || original_rows == 0 { + return 0; + } + if selectivity >= 1.0 { + return original_ndv; + } + + let ndv = original_ndv as f64; + let rows = original_rows as f64; + + let rows_per_value = rows / ndv; + let survival_prob = 1.0 - (1.0 - selectivity).powf(rows_per_value); + let expected_ndv = ndv * survival_prob; + + (expected_ndv.round() as usize).clamp(1, original_ndv) +} + +/// Rescale `total_byte_size` proportionally after overriding `num_rows`. +/// +/// When a provider replaces `num_rows` but keeps the rest of the stats from +/// `partition_statistics`, the original `total_byte_size` becomes inconsistent. +/// This function adjusts it by the ratio `new_rows / old_rows`, preserving the +/// average bytes-per-row from the original estimate. +fn rescale_byte_size(stats: &mut Statistics, new_num_rows: Precision<usize>) { + let old_rows = stats.num_rows; + stats.num_rows = new_num_rows; + stats.total_byte_size = match (old_rows, new_num_rows, stats.total_byte_size) { + (Precision::Exact(old), Precision::Exact(new), Precision::Exact(bytes)) + if old > 0 => + { + Precision::Exact((bytes as f64 * new as f64 / old as f64).round() as usize) + } + _ => match ( + old_rows.get_value(), + new_num_rows.get_value(), + stats.total_byte_size.get_value(), + ) { + (Some(&old), Some(&new), Some(&bytes)) if old > 0 => Precision::Inexact( + (bytes as f64 * new as f64 / old as f64).round() as usize, + ), + _ => stats.total_byte_size, + }, + }; +} + +/// Fetches base statistics from the operator's built-in `partition_statistics`, +/// overrides `num_rows` with the registry-computed estimate, and rescales +/// `total_byte_size` proportionally. +/// +/// Used by providers that compute a better row count but cannot yet propagate +/// column-level stats (NDV, min/max) through the operator — pending #20184. +fn computed_with_row_count( + plan: &dyn ExecutionPlan, + num_rows: Precision<usize>, +) -> Result<StatisticsResult> { + let mut base = Arc::unwrap_or_clone( + StatisticsContext::new().compute(plan, &StatisticsArgs::new())?, + ); + rescale_byte_size(&mut base, num_rows); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(base))) +} + +/// Statistics provider for [`FilterExec`](crate::filter::FilterExec) that uses +/// pre-computed enhanced child statistics from the registry walk. +/// +/// Unlike the default provider (which calls `partition_statistics` and gets raw +/// child stats), this provider receives enhanced child stats that may include +/// NDV overrides injected at the scan level. It applies the same selectivity +/// estimation logic as `FilterExec::statistics_helper`, then additionally +/// adjusts each column's `distinct_count` using [`ndv_after_selectivity`] based +/// on the computed selectivity ratio. +#[derive(Debug, Default)] +pub struct FilterStatisticsProvider; + +impl StatisticsProvider for FilterStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::filter::FilterExec; + + let Some(filter) = plan.downcast_ref::<FilterExec>() else { + return Ok(StatisticsResult::Delegate); + }; + if child_stats.is_empty() { + return Ok(StatisticsResult::Delegate); + } + + let input_stats = (*child_stats[0].base).clone(); + let input_rows = input_stats.num_rows; + let mut stats = FilterExec::statistics_helper( + &filter.input().schema(), + input_stats, + filter.predicate(), + filter.default_selectivity(), + // TODO: pass filter.expression_analyzer_registry() once #21122 lands + )?; + + // Adjust distinct_count for each column using the selectivity ratio + // via the probabilistic survival model from + // ndv_after_selectivity to account for rows removed by the filter. + if let (Some(&orig_rows), Some(&filtered_rows)) = + (input_rows.get_value(), stats.num_rows.get_value()) + && orig_rows > 0 + && filtered_rows < orig_rows + { + let selectivity = filtered_rows as f64 / orig_rows as f64; + for col_stat in &mut stats.column_statistics { + if let Some(&ndv) = col_stat.distinct_count.get_value() { + let adjusted = ndv_after_selectivity(ndv, orig_rows, selectivity); + col_stat.distinct_count = Precision::Inexact(adjusted); + } + } + } + + let stats = stats.project(filter.projection().as_ref()); + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) + } +} + +/// Statistics provider for [`ProjectionExec`](crate::projection::ProjectionExec) +/// that uses pre-computed enhanced child statistics from the registry walk. +/// +/// Maps enhanced child column statistics to output columns based on the +/// projection expressions, preserving NDV and other statistics through +/// column references. +#[derive(Debug, Default)] +pub struct ProjectionStatisticsProvider; + +impl StatisticsProvider for ProjectionStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::projection::ProjectionExec; + + let Some(proj) = plan.downcast_ref::<ProjectionExec>() else { + return Ok(StatisticsResult::Delegate); + }; + if child_stats.is_empty() { + return Ok(StatisticsResult::Delegate); + } + + let input_stats = (*child_stats[0].base).clone(); + let output_schema = proj.schema(); + // TODO: pass proj.expression_analyzer_registry() once #21122 lands, + // so expression-level NDV/min/max feeds into projected column stats. + let stats = proj + .projection_expr() + .project_statistics(input_stats, &output_schema)?; + Ok(StatisticsResult::Computed(ExtendedStatistics::new(stats))) + } +} + +/// Statistics provider for single-input operators with +/// [`CardinalityEffect::Equal`](crate::execution_plan::CardinalityEffect::Equal). +/// +/// These operators (Sort, Repartition, CoalescePartitions, etc.) don't +/// transform statistics, so we pass through the enhanced child stats directly. +/// This avoids the fallback calling `partition_statistics(None)` which would +/// trigger a redundant internal recursion with raw (non-enhanced) stats. +#[derive(Debug, Default)] +pub struct PassthroughStatisticsProvider; + +impl StatisticsProvider for PassthroughStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::execution_plan::CardinalityEffect; + + if child_stats.len() != 1 + || !matches!(plan.cardinality_effect(), CardinalityEffect::Equal) + { + return Ok(StatisticsResult::Delegate); + } + + // Only pass through when the schema is unchanged (same column count). + // Operators like WindowAggExec preserve row count but add columns; + // passing through child stats would produce wrong column_statistics. + let input_cols = child_stats[0].base.column_statistics.len(); + let output_cols = plan.schema().fields().len(); + if input_cols != output_cols { + return Ok(StatisticsResult::Delegate); + } + + Ok(StatisticsResult::Computed(child_stats[0].clone())) + } +} + +/// Statistics provider for [`AggregateExec`](crate::aggregates::AggregateExec) +/// that estimates output cardinality from the NDV of GROUP BY columns. +/// +/// For each GROUP BY column, looks up `distinct_count` from the enhanced +/// child statistics. The estimated output rows is the product of all +/// column NDVs, capped at the input row count. This assumes independence +/// between columns, so correlated columns (e.g., `city` and `state`) will +/// produce overestimates. +/// +/// For GROUPING SETS / CUBE / ROLLUP, delegates to the built-in +/// `partition_statistics`, which handles per-set NDV estimation correctly. +/// +/// Delegates when: +/// - The plan is not an `AggregateExec` +/// - The aggregate is `Partial` (per-partition, not bounded by global NDV) +/// - GROUP BY is empty (scalar aggregate) +/// - Any GROUP BY expression is not a simple column reference +/// - Any GROUP BY column lacks NDV information +#[derive(Debug, Default)] +pub struct AggregateStatisticsProvider; + +impl StatisticsProvider for AggregateStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::aggregates::AggregateExec; + use datafusion_physical_expr::expressions::Column; + + use crate::aggregates::AggregateMode; + + let Some(agg) = plan.downcast_ref::<AggregateExec>() else { + return Ok(StatisticsResult::Delegate); + }; + + // Partial aggregates produce per-partition groups, not bounded by + // global NDV; delegate to the built-in estimate for those. + if matches!(agg.mode(), AggregateMode::Partial) { + return Ok(StatisticsResult::Delegate); + } + + if child_stats.is_empty() || agg.group_expr().expr().is_empty() { + return Ok(StatisticsResult::Delegate); + } + + let input_stats = &child_stats[0].base; + + // Compute NDV product of GROUP BY columns + let mut ndv_product: Option<usize> = None; + for (expr, _) in agg.group_expr().expr().iter() { + let Some(col) = expr.downcast_ref::<Column>() else { + return Ok(StatisticsResult::Delegate); + }; + let Some(&ndv) = input_stats + .column_statistics + .get(col.index()) + .and_then(|s| s.distinct_count.get_value()) + else { + return Ok(StatisticsResult::Delegate); + }; + if ndv == 0 { + return Ok(StatisticsResult::Delegate); + } + ndv_product = Some(match ndv_product { + Some(prev) => prev.saturating_mul(ndv), + None => ndv, + }); + } + + let Some(product) = ndv_product else { + return Ok(StatisticsResult::Delegate); + }; + + // For CUBE/ROLLUP/GROUPING SETS (multiple grouping sets), delegate to + // the built-in estimate, which handles per-set NDV estimation correctly. + if agg.group_expr().groups().len() > 1 { + return Ok(StatisticsResult::Delegate); + } + + // Cap at input rows + let estimate = match input_stats.num_rows.get_value() { + Some(&rows) => product.min(rows), + None => product, + }; + + let num_rows = Precision::Inexact(estimate); + + computed_with_row_count(plan, num_rows) + } +} + +/// Statistics provider for equi-joins (hash join, sort-merge join) and cross joins. +/// +/// For equi-joins, estimates output cardinality as +/// `left_rows * right_rows / product(max(left_ndv_i, right_ndv_i))` +/// across all join key columns (assuming independence between keys), +/// falling back to the Cartesian product when any key lacks NDV on both sides. +/// For cross joins, uses the exact Cartesian product. +/// +/// The base inner-join estimate is then adjusted for the join type: +/// - Semi joins: capped at the preserved-side row count +/// - Anti joins: preserved-side minus matched rows (clamped to 0) +/// - Left/Right outer: at least as many rows as the preserved side +/// - Full outer: at least `left + right - inner_estimate` +/// - Left mark: exactly `left_rows` (one output row per left row) +/// +/// Delegates when: +/// - The plan is not a supported join type +/// - Either input lacks row count information +#[derive(Debug, Default)] +pub struct JoinStatisticsProvider; + +impl StatisticsProvider for JoinStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::joins::{CrossJoinExec, HashJoinExec, SortMergeJoinExec}; + use datafusion_common::JoinType; + use datafusion_physical_expr::expressions::Column; + + if child_stats.len() < 2 { + return Ok(StatisticsResult::Delegate); + } + + let left = &child_stats[0].base; + let right = &child_stats[1].base; + + let (Some(&left_rows), Some(&right_rows)) = + (left.num_rows.get_value(), right.num_rows.get_value()) + else { + return Ok(StatisticsResult::Delegate); + }; + + use crate::joins::JoinOnRef; + + /// Estimate equi-join output using NDV of join key columns: + /// left_rows * right_rows / product(max(left_ndv_i, right_ndv_i)) + /// Falls back to Cartesian product if any key lacks NDV on both sides. + fn equi_join_estimate( + on: JoinOnRef, + left: &Statistics, + right: &Statistics, + left_rows: usize, + right_rows: usize, + ) -> usize { + if on.is_empty() { + return left_rows.saturating_mul(right_rows); + } + let mut ndv_divisor: usize = 1; + for (left_key, right_key) in on { + let left_ndv = left_key + .downcast_ref::<Column>() + .and_then(|c| left.column_statistics.get(c.index())) + .and_then(|s| s.distinct_count.get_value().copied()); + let right_ndv = right_key + .downcast_ref::<Column>() + .and_then(|c| right.column_statistics.get(c.index())) + .and_then(|s| s.distinct_count.get_value().copied()); + match (left_ndv, right_ndv) { + (Some(l), Some(r)) if l > 0 && r > 0 => { + ndv_divisor = ndv_divisor.saturating_mul(l.max(r)); + } + _ => return left_rows.saturating_mul(right_rows), + } + } + let max_rows = left_rows.saturating_mul(right_rows); + max_rows.checked_div(ndv_divisor).unwrap_or(max_rows) + } + + let (inner_estimate, is_exact_cartesian, join_type) = if let Some(hash_join) = + plan.downcast_ref::<HashJoinExec>() + { + let est = + equi_join_estimate(hash_join.on(), left, right, left_rows, right_rows); + (est, false, *hash_join.join_type()) + } else if let Some(smj) = plan.downcast_ref::<SortMergeJoinExec>() { + let est = equi_join_estimate(smj.on(), left, right, left_rows, right_rows); + (est, false, smj.join_type()) + } else if plan.downcast_ref::<CrossJoinExec>().is_some() { + let both_exact = left.num_rows.is_exact().unwrap_or(false) + && right.num_rows.is_exact().unwrap_or(false); + ( + left_rows.saturating_mul(right_rows), + both_exact, + JoinType::Inner, + ) + } else { + return Ok(StatisticsResult::Delegate); + }; + + // Apply join-type-aware cardinality bounds + let estimated = match join_type { + JoinType::Inner => inner_estimate, + JoinType::Left => inner_estimate.max(left_rows), + JoinType::Right => inner_estimate.max(right_rows), + JoinType::Full => { + // At least left + right - matched, but never less than inner + let outer_bound = left_rows + .saturating_add(right_rows) + .saturating_sub(inner_estimate); + inner_estimate.max(outer_bound) + } + JoinType::LeftSemi => inner_estimate.min(left_rows), + JoinType::RightSemi => inner_estimate.min(right_rows), + JoinType::LeftAnti => left_rows.saturating_sub(inner_estimate.min(left_rows)), + JoinType::RightAnti => { + right_rows.saturating_sub(inner_estimate.min(right_rows)) + } + JoinType::LeftMark => left_rows, + JoinType::RightMark => right_rows, + }; + + // NL join inner with exact inputs is an exact Cartesian product; + // NDV-based estimates are inherently inexact. + let num_rows = if is_exact_cartesian && join_type == JoinType::Inner { + Precision::Exact(estimated) + } else { + Precision::Inexact(estimated) + }; + + computed_with_row_count(plan, num_rows) + } +} + +/// Statistics provider for [`LocalLimitExec`](crate::limit::LocalLimitExec) and +/// [`GlobalLimitExec`](crate::limit::GlobalLimitExec). +/// +/// Caps output row count at the limit value, accounting for any leading skip offset +/// in `GlobalLimitExec`. +#[derive(Debug, Default)] +pub struct LimitStatisticsProvider; + +impl StatisticsProvider for LimitStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::limit::{GlobalLimitExec, LocalLimitExec}; + + if child_stats.is_empty() { + return Ok(StatisticsResult::Delegate); + } + + let (skip, fetch) = if let Some(limit) = plan.downcast_ref::<LocalLimitExec>() { + (0usize, Some(limit.fetch())) + } else if let Some(limit) = plan.downcast_ref::<GlobalLimitExec>() { + (limit.skip(), limit.fetch()) + } else { + return Ok(StatisticsResult::Delegate); + }; + + let num_rows = match child_stats[0].base.num_rows { + Precision::Exact(rows) => { + let available = rows.saturating_sub(skip); + Precision::Exact(fetch.map_or(available, |f| available.min(f))) + } + Precision::Inexact(rows) => { + let available = rows.saturating_sub(skip); + match fetch { + Some(f) => Precision::Inexact(available.min(f)), + None => Precision::Inexact(available), + } + } + Precision::Absent => match fetch { + Some(f) => Precision::Inexact(f), + None => Precision::Absent, + }, + }; + + computed_with_row_count(plan, num_rows) + } +} + +/// Statistics provider for [`UnionExec`](crate::union::UnionExec). +/// +/// Sums row counts across all inputs. +#[derive(Debug, Default)] +pub struct UnionStatisticsProvider; + +impl StatisticsProvider for UnionStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + use crate::union::UnionExec; + + if plan.downcast_ref::<UnionExec>().is_none() { + return Ok(StatisticsResult::Delegate); + } + + let total = child_stats.iter().try_fold( + Precision::Exact(0usize), + |acc, s| -> Result<Precision<usize>> { + Ok(match (acc, s.base.num_rows) { + (Precision::Absent, _) | (_, Precision::Absent) => Precision::Absent, + (Precision::Exact(a), Precision::Exact(b)) => { + Precision::Exact(a.saturating_add(b)) + } + (Precision::Inexact(a), Precision::Exact(b)) + | (Precision::Exact(a), Precision::Inexact(b)) + | (Precision::Inexact(a), Precision::Inexact(b)) => { + Precision::Inexact(a.saturating_add(b)) + } + }) + }, + )?; + + computed_with_row_count(plan, total) + } +} + +type ProviderFn = dyn Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result<StatisticsResult> + + Send + + Sync; + +/// A [`StatisticsProvider`] backed by a user-supplied closure. +/// +/// Useful for injecting custom statistics in tests or for cardinality feedback +/// pipelines where real runtime statistics need to override plan estimates. +/// The closure receives the current plan node and its children's enhanced +/// statistics, returning a [`StatisticsResult`]. +/// +/// To distinguish between multiple nodes of the same type (e.g., two +/// `FilterExec` nodes), match on structural properties like the input schema's +/// column names, number of columns, or child row counts. +/// +/// # Example +/// +/// ```rust,ignore (requires crate-internal imports) +/// let provider = ClosureStatisticsProvider::new(|plan, child_stats| { +/// if plan.downcast_ref::<FilterExec>().is_some() { +/// Ok(StatisticsResult::Computed(ExtendedStatistics::from(Statistics { +/// num_rows: Precision::Inexact(42), +/// ..Statistics::new_unknown(plan.schema().as_ref()) +/// }))) +/// } else { +/// Ok(StatisticsResult::Delegate) +/// } +/// }); +/// ``` +pub struct ClosureStatisticsProvider { + f: Box<ProviderFn>, +} + +impl ClosureStatisticsProvider { + /// Create a new provider from a closure. + pub fn new( + f: impl Fn(&dyn ExecutionPlan, &[ExtendedStatistics]) -> Result<StatisticsResult> + + Send + + Sync + + 'static, + ) -> Self { + Self { f: Box::new(f) } + } +} + +impl Debug for ClosureStatisticsProvider { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "ClosureStatisticsProvider") + } +} + +impl StatisticsProvider for ClosureStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + (self.f)(plan, child_stats) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::filter::FilterExec; + use crate::projection::ProjectionExec; + use crate::statistics::StatisticsArgs; + use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, PlanProperties, + ReplaceChildrenOptions, + }; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::stats::Precision; + use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, col, lit}; + use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; + use std::fmt; + + use crate::execution_plan::{Boundedness, EmissionType}; + use datafusion_common::tree_node::TreeNodeRecursion; + + fn make_schema() -> Arc<Schema> { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])) + } + + #[derive(Debug)] + struct MockSourceExec { + schema: Arc<Schema>, + stats: Statistics, + cache: Arc<PlanProperties>, + } + + impl MockSourceExec { + fn new(schema: Arc<Schema>, num_rows: Precision<usize>) -> Self { + let num_cols = schema.fields().len(); + Self::with_column_stats( + schema, + num_rows, + vec![ColumnStatistics::new_unknown(); num_cols], + ) + } + + fn with_column_stats( + schema: Arc<Schema>, + num_rows: Precision<usize>, + column_statistics: Vec<ColumnStatistics>, + ) -> Self { + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + let cache = Arc::new(PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + schema, + stats: Statistics { + num_rows, + total_byte_size: Precision::Absent, + column_statistics, + }, + cache, + } + } + } + + impl DisplayAs for MockSourceExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "MockSourceExec") + } + } + + impl ExecutionPlan for MockSourceExec { + fn name(&self) -> &str { + "MockSourceExec" + } + + fn schema(&self) -> Arc<Schema> { + Arc::clone(&self.schema) + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<datafusion_execution::TaskContext>, + ) -> Result<crate::SendableRecordBatchStream> { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::new(self.stats.clone())) + } + } + + fn make_source(num_rows: usize) -> Arc<dyn ExecutionPlan> { + Arc::new(MockSourceExec::new( + make_schema(), + Precision::Exact(num_rows), + )) + } + + #[test] + fn test_default_provider() -> Result<()> { + let engine = StatisticsRegistry::new(); + let source = make_source(1000); + + let stats = engine.compute(source.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); + Ok(()) + } + + #[test] + fn test_custom_chain_configuration() -> Result<()> { + let source = make_source(1000); + + // Test with_providers: fully custom chain (no default) + let custom_only = + StatisticsRegistry::with_providers(vec![Arc::new(CustomStatisticsProvider)]); + // CustomStatisticsProvider only handles CustomExec, delegates for others + // With no default provider, filter returns fallback statistics + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(lit(true), Arc::clone(&source))?); + let stats = custom_only.compute(filter.as_ref())?; + // Falls back to plan.statistics() since no provider handles it + assert!(stats.base.num_rows.get_value().is_some()); + + // Test with_providers: custom provider + built-in fallback + let with_override = + StatisticsRegistry::with_providers(vec![Arc::new(OverrideFilterProvider { + fixed_selectivity: 0.25, + }) + as Arc<dyn StatisticsProvider>]); + // OverrideFilterProvider handles filters, built-in fallback handles the rest + let stats = with_override.compute(filter.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Inexact(250))); + + // Verify chain inspection + assert_eq!(with_override.providers().len(), 1); + + Ok(()) + } + + #[derive(Debug)] + struct CustomExec { + input: Arc<dyn ExecutionPlan>, + } + + impl DisplayAs for CustomExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "CustomExec") + } + } + + impl ExecutionPlan for CustomExec { + fn name(&self) -> &str { + "CustomExec" + } + + fn schema(&self) -> Arc<Schema> { + self.input.schema() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(CustomExec { + input: Arc::clone(&children[0]), + })) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn properties(&self) -> &Arc<PlanProperties> { + self.input.properties() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<datafusion_execution::TaskContext>, + ) -> Result<crate::SendableRecordBatchStream> { + unimplemented!() + } + } + + #[derive(Debug)] + struct CustomStatisticsProvider; + + impl StatisticsProvider for CustomStatisticsProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + if plan.downcast_ref::<CustomExec>().is_some() { + Ok(StatisticsResult::Computed(child_stats[0].clone())) + } else { + Ok(StatisticsResult::Delegate) + } + } + } + + #[test] + fn test_custom_provider_for_custom_exec() -> Result<()> { + let mut engine = StatisticsRegistry::new(); + engine.register(Arc::new(CustomStatisticsProvider)); + + let source = make_source(1000); + let custom: Arc<dyn ExecutionPlan> = Arc::new(CustomExec { input: source }); + + let stats = engine.compute(custom.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); + Ok(()) + } + + #[derive(Debug)] + struct OverrideFilterProvider { + fixed_selectivity: f64, + } + + impl StatisticsProvider for OverrideFilterProvider { + fn compute_statistics( + &self, + plan: &dyn ExecutionPlan, + child_stats: &[ExtendedStatistics], + ) -> Result<StatisticsResult> { + if plan.downcast_ref::<FilterExec>().is_some() { + if let Some(&input_rows) = child_stats[0].base.num_rows.get_value() { + let estimated = (input_rows as f64 * self.fixed_selectivity) as usize; + Ok(StatisticsResult::Computed(ExtendedStatistics::from( + Statistics { + num_rows: Precision::Inexact(estimated), + total_byte_size: Precision::Absent, + column_statistics: child_stats[0] + .base + .column_statistics + .clone(), + }, + ))) + } else { + Ok(StatisticsResult::Delegate) + } + } else { + Ok(StatisticsResult::Delegate) + } + } + } + + #[test] + fn test_override_builtin_operator() -> Result<()> { + let mut engine = StatisticsRegistry::new(); + engine.register(Arc::new(OverrideFilterProvider { + fixed_selectivity: 0.1, + })); + + let source = make_source(1000); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(lit(true), source)?); + + let stats = engine.compute(filter.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Inexact(100))); + Ok(()) + } + + #[test] + fn test_filter_statistics_propagation() -> Result<()> { + let engine = StatisticsRegistry::new(); + let source = make_source(1000); + let predicate = lit(true); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, source)?); + + let stats = engine.compute(filter.as_ref())?; + assert!(stats.base.num_rows.get_value().unwrap_or(&0) <= &1000); + Ok(()) + } + + #[test] + fn test_filter_adjusts_ndv_by_selectivity() -> Result<()> { + use datafusion_common::ScalarValue; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{ + BinaryExpr, Column as PhysColumn, Literal, + }; + + // Source: 1000 rows, NDV(a)=1000 (unique), NDV(b)=800 (near-unique) + // With NDV close to num_rows, each value has ~1.25 rows, so filtering + // visibly reduces the number of surviving distinct values. + let schema = make_schema(); // "a" Int32, "b" Int32 + let col_stats = vec![ + { + let mut cs = ColumnStatistics::new_unknown(); + cs.distinct_count = Precision::Exact(1000); + cs.min_value = Precision::Exact(ScalarValue::Int32(Some(1))); + cs.max_value = Precision::Exact(ScalarValue::Int32(Some(1000))); + cs + }, + { + let mut cs = ColumnStatistics::new_unknown(); + cs.distinct_count = Precision::Exact(800); + cs.min_value = Precision::Exact(ScalarValue::Int32(Some(1))); + cs.max_value = Precision::Exact(ScalarValue::Int32(Some(800))); + cs + }, + ]; + let source: Arc<dyn ExecutionPlan> = Arc::new(MockSourceExec::with_column_stats( + schema, + Precision::Exact(1000), + col_stats, + )); + + // Filter: a > 900 (selectivity ~10%, keeps values 901-1000) + let predicate: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new( + Arc::new(PhysColumn::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(900)))), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, source)?); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(FilterStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(filter.as_ref())?; + + let output_ndv_a = stats.base.column_statistics[0] + .distinct_count + .get_value() + .copied() + .unwrap_or(0); + let output_ndv_b = stats.base.column_statistics[1] + .distinct_count + .get_value() + .copied() + .unwrap_or(0); + + // NDV(a): interval analysis narrows to [901,1000] -> ~100 distinct values + assert!( + output_ndv_a <= 100, + "Expected NDV(a) <= 100 after filter, got {output_ndv_a}" + ); + // NDV(b): not in predicate, but selectivity ~10% with 1.25 rows/value + // means many distinct values are lost. ndv_after_selectivity(800, 1000, 0.1) + // gives ~76. Significantly less than the original 800. + assert!( + output_ndv_b < 200, + "Expected NDV(b) < 200 after filter, got {output_ndv_b}" + ); + Ok(()) + } + + #[test] + fn test_projection_statistics_propagation() -> Result<()> { + let engine = StatisticsRegistry::new(); + let source = make_source(1000); + let schema = make_schema(); + let proj: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new( + vec![(col("a", &schema)?, "a".to_string())], + source, + )?); + + let stats = engine.compute(proj.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); + Ok(()) + } + + #[test] + fn test_passthrough_statistics_propagation() -> Result<()> { + use crate::coalesce_partitions::CoalescePartitionsExec; + + let engine = StatisticsRegistry::new(); + let source = make_source(1000); + let coalesce: Arc<dyn ExecutionPlan> = + Arc::new(CoalescePartitionsExec::new(source)); + + let stats = engine.compute(coalesce.as_ref())?; + // PassthroughStatisticsProvider should propagate child row count unchanged + assert_eq!(stats.base.num_rows, Precision::Exact(1000)); + Ok(()) + } + + #[test] + fn test_chain_priority() -> Result<()> { + let mut engine = StatisticsRegistry::new(); + engine.register(Arc::new(OverrideFilterProvider { + fixed_selectivity: 0.5, + })); + engine.register(Arc::new(CustomStatisticsProvider)); + + let source = make_source(1000); + + // CustomExec handled by CustomStatisticsProvider + let custom: Arc<dyn ExecutionPlan> = Arc::new(CustomExec { + input: Arc::clone(&source), + }); + let stats = engine.compute(custom.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Exact(1000))); + + // FilterExec: CustomStatisticsProvider delegates, OverrideFilterProvider handles + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(lit(true), source)?); + let stats = engine.compute(filter.as_ref())?; + assert!(matches!(stats.base.num_rows, Precision::Inexact(500))); + + Ok(()) + } + + // ========================================================================= + // num_distinct_vals Utility Tests + // ========================================================================= + + #[test] + fn test_num_distinct_vals_basic() { + assert_eq!(num_distinct_vals(0, 100), 0); + assert_eq!(num_distinct_vals(100, 0), 0); + assert_eq!(num_distinct_vals(100, 100), 100); + assert_eq!(num_distinct_vals(100, 200), 100); + + let ndv = num_distinct_vals(1000, 100); + assert!((90..=100).contains(&ndv), "Expected ~95, got {ndv}"); + + let ndv = num_distinct_vals(1000, 500); + assert!((350..=450).contains(&ndv), "Expected ~393, got {ndv}"); + + let ndv = num_distinct_vals(1_000_000, 10_000); + assert!((9900..=10000).contains(&ndv), "Expected ~9950, got {ndv}"); + + let ndv = num_distinct_vals(1_000_000, 100); + assert!((99..=100).contains(&ndv), "Expected ~100, got {ndv}"); + } + + #[test] + fn test_num_distinct_vals_small_domain() { + let ndv = num_distinct_vals(10, 5); + assert!((3..=5).contains(&ndv), "Expected ~4, got {ndv}"); + + assert_eq!(num_distinct_vals(10, 20), 10); + assert_eq!(num_distinct_vals(10, 1), 1); + } + + #[test] + fn test_ndv_after_selectivity() { + let ndv = ndv_after_selectivity(1000, 10000, 0.1); + assert!((600..=700).contains(&ndv), "Expected ~632, got {ndv}"); + + let ndv = ndv_after_selectivity(1000, 10000, 0.01); + assert!((90..=100).contains(&ndv), "Expected ~95, got {ndv}"); + + assert_eq!(ndv_after_selectivity(1000, 10000, 0.0), 0); + assert_eq!(ndv_after_selectivity(1000, 10000, 1.0), 1000); + assert_eq!(ndv_after_selectivity(0, 10000, 0.5), 0); + } + + // ========================================================================= + // AggregateStatisticsProvider tests + // ========================================================================= + + use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + + fn make_source_with_ndv( + num_rows: usize, + col_ndvs: Vec<Option<usize>>, + ) -> Arc<dyn ExecutionPlan> { + let fields: Vec<Field> = col_ndvs + .iter() + .enumerate() + .map(|(i, _)| Field::new(format!("c{i}"), DataType::Int32, false)) + .collect(); + let schema = Arc::new(Schema::new(fields)); + let col_stats = col_ndvs + .into_iter() + .map(|ndv| { + let mut cs = ColumnStatistics::new_unknown(); + if let Some(n) = ndv { + cs.distinct_count = Precision::Exact(n); + } + cs + }) + .collect(); + Arc::new(MockSourceExec::with_column_stats( + schema, + Precision::Exact(num_rows), + col_stats, + )) + } + + fn make_aggregate( + input: Arc<dyn ExecutionPlan>, + group_by: PhysicalGroupBy, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![], + vec![], + Arc::clone(&input), + input.schema(), + )?)) + } + + #[test] + fn test_aggregate_provider_with_ndv() -> Result<()> { + let source = make_source_with_ndv(100, vec![Some(10)]); + let group_by = PhysicalGroupBy::new_single(vec![( + Arc::new(Column::new("c0", 0)), + "c0".to_string(), + )]); + let agg = make_aggregate(source, group_by)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(10)); + Ok(()) + } + + #[test] + fn test_aggregate_provider_multi_column() -> Result<()> { + let source = make_source_with_ndv(1000, vec![Some(10), Some(5)]); + let group_by = PhysicalGroupBy::new_single(vec![ + (Arc::new(Column::new("c0", 0)), "c0".to_string()), + (Arc::new(Column::new("c1", 1)), "c1".to_string()), + ]); + let agg = make_aggregate(source, group_by)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + // 10 * 5 = 50 + assert_eq!(stats.base.num_rows, Precision::Inexact(50)); + Ok(()) + } + + #[test] + fn test_aggregate_provider_caps_at_input_rows() -> Result<()> { + // NDV product (100 * 100 = 10_000) exceeds input rows (500) + let source = make_source_with_ndv(500, vec![Some(100), Some(100)]); + let group_by = PhysicalGroupBy::new_single(vec![ + (Arc::new(Column::new("c0", 0)), "c0".to_string()), + (Arc::new(Column::new("c1", 1)), "c1".to_string()), + ]); + let agg = make_aggregate(source, group_by)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(500)); + Ok(()) + } + + #[test] + fn test_aggregate_provider_no_ndv_delegates() -> Result<()> { + // No NDV on the GROUP BY column + let source = make_source_with_ndv(100, vec![None]); + let group_by = PhysicalGroupBy::new_single(vec![( + Arc::new(Column::new("c0", 0)), + "c0".to_string(), + )]); + let agg = make_aggregate(source, group_by)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + // Delegates to DefaultStatisticsProvider, which calls partition_statistics + assert!( + stats.base.num_rows.get_value().is_some() + || matches!(stats.base.num_rows, Precision::Absent) + ); + Ok(()) + } + + #[test] + fn test_aggregate_provider_non_column_expr_delegates() -> Result<()> { + let source = make_source_with_ndv(100, vec![Some(10), Some(5)]); + // GROUP BY an expression (c0 + c1), not a simple column ref + let expr: Arc<dyn PhysicalExpr> = Arc::new(BinaryExpr::new( + Arc::new(Column::new("c0", 0)), + Operator::Plus, + Arc::new(Column::new("c1", 1)), + )); + let group_by = PhysicalGroupBy::new_single(vec![(expr, "sum".to_string())]); + let agg = make_aggregate(source, group_by)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + // Should delegate (expression is not a Column) + assert!( + stats.base.num_rows.get_value().is_some() + || matches!(stats.base.num_rows, Precision::Absent) + ); + Ok(()) + } + + #[test] + fn test_aggregate_provider_grouping_sets() -> Result<()> { + let source = make_source_with_ndv(1000, vec![Some(10), Some(5)]); + // GROUPING SETS: (c0, c1), (c0), (c1) -> 3 groups + let group_by = PhysicalGroupBy::new( + vec![ + (Arc::new(Column::new("c0", 0)), "c0".to_string()), + (Arc::new(Column::new("c1", 1)), "c1".to_string()), + ], + vec![ + ( + Arc::new(Literal::new(ScalarValue::Int32(None))), + "c0".to_string(), + ), + ( + Arc::new(Literal::new(ScalarValue::Int32(None))), + "c1".to_string(), + ), + ], + vec![ + vec![false, true], // (c0, NULL) - group by c0 only + vec![true, false], // (NULL, c1) - group by c1 only + vec![false, false], // (c0, c1) - group by both + ], + true, + ); + let agg = make_aggregate(source, group_by)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + // Multiple grouping sets: provider delegates to DefaultStatisticsProvider, + // which calls the built-in partition_statistics for correct per-set + // NDV estimation. The exact value depends on the built-in implementation. + assert!( + stats.base.num_rows.get_value().is_some() + || matches!(stats.base.num_rows, Precision::Absent) + ); + Ok(()) + } + + #[test] + fn test_aggregate_provider_partial_delegates() -> Result<()> { + // Partial aggregates produce per-partition groups; the provider + // should delegate rather than applying global NDV bounds. + let source = make_source_with_ndv(100, vec![Some(10)]); + let group_by = PhysicalGroupBy::new_single(vec![( + Arc::new(Column::new("c0", 0)), + "c0".to_string(), + )]); + let agg: Arc<dyn ExecutionPlan> = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + vec![], + vec![], + Arc::clone(&source), + source.schema(), + )?); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(AggregateStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(agg.as_ref())?; + // Should fall through to DefaultStatisticsProvider (partition_statistics). + // The exact value depends on the built-in implementation. + assert!( + stats.base.num_rows.get_value().is_some() + || matches!(stats.base.num_rows, Precision::Absent) + ); + Ok(()) + } + + // ========================================================================= + // JoinStatisticsProvider tests + // ========================================================================= + + use crate::joins::{HashJoinExec, PartitionMode}; + use datafusion_common::{JoinType, NullEquality}; + + fn make_source_with_ndv_2col( + num_rows: usize, + ndv_a: Option<usize>, + ) -> Arc<dyn ExecutionPlan> { + let schema = make_schema(); // "a" Int32, "b" Int32 + let col_stats = vec![ + { + let mut cs = ColumnStatistics::new_unknown(); + if let Some(n) = ndv_a { + cs.distinct_count = Precision::Exact(n); + } + cs + }, + ColumnStatistics::new_unknown(), + ]; + Arc::new(MockSourceExec::with_column_stats( + schema, + Precision::Exact(num_rows), + col_stats, + )) + } + + fn make_hash_join( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let _schema = make_schema(); + let on: crate::joins::JoinOn = vec![( + Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>, + Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>, + )]; + Ok(Arc::new(HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?)) + } + + #[test] + fn test_join_provider_with_ndv() -> Result<()> { + // left: 1000 rows, NDV(a)=100; right: 500 rows, NDV(a)=50 + // expected = 1000 * 500 / max(100, 50) = 5000 + let left = make_source_with_ndv_2col(1000, Some(100)); + let right = make_source_with_ndv_2col(500, Some(50)); + let join = make_hash_join(left, right)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(join.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(5000)); + Ok(()) + } + + #[test] + fn test_join_provider_uses_actual_key_column_ndv() -> Result<()> { + // Join on column "b" (index 1), NDV only set on "b", not "a". + // Old first()-based code would look up column 0 (a), find no NDV, + // and fall back to Cartesian product. The fix looks up column 1 (b). + // left: 1000 rows, NDV(b)=50; right: 500 rows, NDV(b)=25 + // expected = 1000 * 500 / max(50, 25) = 10000 + let schema = make_schema(); // "a" Int32, "b" Int32 + let make_source_ndv_b = + |num_rows: usize, ndv_b: usize| -> Arc<dyn ExecutionPlan> { + let col_stats = vec![ + ColumnStatistics::new_unknown(), // "a": no NDV + { + let mut cs = ColumnStatistics::new_unknown(); + cs.distinct_count = Precision::Exact(ndv_b); + cs + }, + ]; + Arc::new(MockSourceExec::with_column_stats( + Arc::clone(&schema), + Precision::Exact(num_rows), + col_stats, + )) + }; + + let left = make_source_ndv_b(1000, 50); + let right = make_source_ndv_b(500, 25); + + // Join on column "b" (index 1) + let on: crate::joins::JoinOn = vec![( + Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>, + Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>, + )]; + let join: Arc<dyn ExecutionPlan> = Arc::new(HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(join.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(10_000)); + Ok(()) + } + + #[test] + fn test_join_provider_multi_key_ndv() -> Result<()> { + // Multi-key join: ON a.a = b.a AND a.b = b.b + // left: 1000 rows, NDV(a)=100, NDV(b)=20 + // right: 500 rows, NDV(a)=50, NDV(b)=10 + // expected = 1000 * 500 / (max(100,50) * max(20,10)) = 500000 / 2000 = 250 + let schema = make_schema(); // "a" Int32, "b" Int32 + let make_source_2ndv = + |num_rows: usize, ndv_a: usize, ndv_b: usize| -> Arc<dyn ExecutionPlan> { + let col_stats = vec![ + { + let mut cs = ColumnStatistics::new_unknown(); + cs.distinct_count = Precision::Exact(ndv_a); + cs + }, + { + let mut cs = ColumnStatistics::new_unknown(); + cs.distinct_count = Precision::Exact(ndv_b); + cs + }, + ]; + Arc::new(MockSourceExec::with_column_stats( + Arc::clone(&schema), + Precision::Exact(num_rows), + col_stats, + )) + }; + + let left = make_source_2ndv(1000, 100, 20); + let right = make_source_2ndv(500, 50, 10); + + let on: crate::joins::JoinOn = vec![ + ( + Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>, + Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>, + ), + ( + Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>, + Arc::new(Column::new("b", 1)) as Arc<dyn PhysicalExpr>, + ), + ]; + let join: Arc<dyn ExecutionPlan> = Arc::new(HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(join.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(250)); + Ok(()) + } + + #[test] + fn test_join_provider_fallback_cartesian() -> Result<()> { + // No NDV available -> Cartesian product estimate + let left = make_source_with_ndv_2col(100, None); + let right = make_source_with_ndv_2col(200, None); + let join = make_hash_join(left, right)?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(join.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(20_000)); + Ok(()) + } + + #[test] + fn test_nl_join_delegates() -> Result<()> { + use crate::joins::NestedLoopJoinExec; + + // NL join delegates to the built-in (NestedLoopJoinExec may have an + // arbitrary JoinFilter, so the provider cannot safely assume Cartesian). + let left = make_source(100); + let right = make_source(200); + let join: Arc<dyn ExecutionPlan> = Arc::new(NestedLoopJoinExec::try_new( + left, + right, + None, + &JoinType::Inner, + None, + )?); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(join.as_ref())?; + // Provider delegates; result comes from built-in partition_statistics. + assert!( + stats.base.num_rows.get_value().is_some() + || matches!(stats.base.num_rows, Precision::Absent) + ); + Ok(()) + } + + fn make_hash_join_typed( + left: Arc<dyn ExecutionPlan>, + right: Arc<dyn ExecutionPlan>, + join_type: JoinType, + ) -> Result<Arc<dyn ExecutionPlan>> { + let on: crate::joins::JoinOn = vec![( + Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>, + Arc::new(Column::new("a", 0)) as Arc<dyn PhysicalExpr>, + )]; + Ok(Arc::new(HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?)) + } + + fn compute_join_rows( + left_rows: usize, + left_ndv: Option<usize>, + right_rows: usize, + right_ndv: Option<usize>, + join_type: JoinType, + ) -> Result<Precision<usize>> { + let left = make_source_with_ndv_2col(left_rows, left_ndv); + let right = make_source_with_ndv_2col(right_rows, right_ndv); + let join = make_hash_join_typed(left, right, join_type)?; + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + Ok(registry.compute(join.as_ref())?.base.num_rows) + } + + #[test] + fn test_join_provider_left_outer() -> Result<()> { + // left=1000, right=500, NDV(a)=100/50 + // inner estimate = 1000*500/100 = 5000, already >= left_rows + // Left outer: max(5000, 1000) = 5000 + assert_eq!( + compute_join_rows(1000, Some(100), 500, Some(50), JoinType::Left)?, + Precision::Inexact(5000) + ); + // Small inner estimate: left=1000, right=10, NDV=100/100 + // inner = 1000*10/100 = 100, left outer = max(100, 1000) = 1000 + assert_eq!( + compute_join_rows(1000, Some(100), 10, Some(100), JoinType::Left)?, + Precision::Inexact(1000) + ); + Ok(()) + } + + #[test] + fn test_join_provider_right_outer() -> Result<()> { + // inner = 1000*10/100 = 100, right outer = max(100, 10) = 100 + assert_eq!( + compute_join_rows(1000, Some(100), 10, Some(100), JoinType::Right)?, + Precision::Inexact(100) + ); + // inner = 10*1000/100 = 100, right outer = max(100, 1000) = 1000 + assert_eq!( + compute_join_rows(10, Some(100), 1000, Some(100), JoinType::Right)?, + Precision::Inexact(1000) + ); + Ok(()) + } + + #[test] + fn test_join_provider_semi_join() -> Result<()> { + // inner = 5000, left semi = min(5000, 1000) = 1000 + assert_eq!( + compute_join_rows(1000, Some(100), 500, Some(50), JoinType::LeftSemi)?, + Precision::Inexact(1000) + ); + // inner = 5000, right semi = min(5000, 500) = 500 + assert_eq!( + compute_join_rows(1000, Some(100), 500, Some(50), JoinType::RightSemi)?, + Precision::Inexact(500) + ); + // Cartesian fallback (no NDV): inner = 1000*500 = 500000, + // left semi = min(500000, 1000) = 1000 (selectivity = 1.0) + assert_eq!( + compute_join_rows(1000, None, 500, None, JoinType::LeftSemi)?, + Precision::Inexact(1000) + ); + Ok(()) + } + + #[test] + fn test_join_provider_anti_join() -> Result<()> { + // inner = 1000*10/100 = 100, left anti = 1000 - min(100, 1000) = 900 + assert_eq!( + compute_join_rows(1000, Some(100), 10, Some(100), JoinType::LeftAnti)?, + Precision::Inexact(900) + ); + // inner = 5000, right anti = 500 - min(5000, 500) = 0 + assert_eq!( + compute_join_rows(1000, Some(100), 500, Some(50), JoinType::RightAnti)?, + Precision::Inexact(0) + ); + Ok(()) + } + + // ========================================================================= + // CrossJoinExec tests (handled by JoinStatisticsProvider) + // ========================================================================= + + #[test] + fn test_cross_join_provider_exact() -> Result<()> { + use crate::joins::CrossJoinExec; + let left = make_source(100); + let right = make_source(200); + let join: Arc<dyn ExecutionPlan> = Arc::new(CrossJoinExec::new(left, right)); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(JoinStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(join.as_ref())?; + // Both inputs have Exact row counts -> result is also Exact + assert_eq!(stats.base.num_rows, Precision::Exact(20_000)); + Ok(()) + } + + // ========================================================================= + // LimitStatisticsProvider tests + // ========================================================================= + + use crate::limit::{GlobalLimitExec, LocalLimitExec}; + + #[test] + fn test_limit_provider_caps_output() -> Result<()> { + // input > fetch -> capped at fetch + let source = make_source(1000); + let limit: Arc<dyn ExecutionPlan> = Arc::new(LocalLimitExec::new(source, 100)); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(LimitStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(limit.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Exact(100)); + Ok(()) + } + + #[test] + fn test_limit_provider_input_smaller_than_fetch() -> Result<()> { + // input < fetch -> output = input + let source = make_source(50); + let limit: Arc<dyn ExecutionPlan> = Arc::new(LocalLimitExec::new(source, 200)); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(LimitStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(limit.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Exact(50)); + Ok(()) + } + + #[test] + fn test_global_limit_provider_skip_and_fetch() -> Result<()> { + // 1000 rows, skip 200, fetch 100 -> exactly 100 + let source = make_source(1000); + let limit: Arc<dyn ExecutionPlan> = + Arc::new(GlobalLimitExec::new(source, 200, Some(100))); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(LimitStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(limit.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Exact(100)); + Ok(()) + } + + #[test] + fn test_global_limit_provider_skip_exceeds_rows() -> Result<()> { + // 100 rows, skip 200 -> 0 rows (skip > available) + let source = make_source(100); + let limit: Arc<dyn ExecutionPlan> = + Arc::new(GlobalLimitExec::new(source, 200, Some(50))); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(LimitStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(limit.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Exact(0)); + Ok(()) + } + + #[test] + fn test_limit_provider_inexact_input() -> Result<()> { + // Inexact(1000) with fetch=100: result must stay Inexact, not Exact, + // because the actual row count could be less than 100. + let source = make_source_with_precision(Precision::Inexact(1000)); + let limit: Arc<dyn ExecutionPlan> = Arc::new(LocalLimitExec::new(source, 100)); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(LimitStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(limit.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(100)); + Ok(()) + } + + // ========================================================================= + // UnionStatisticsProvider tests + // ========================================================================= + + use crate::union::UnionExec; + + fn make_source_with_precision(num_rows: Precision<usize>) -> Arc<dyn ExecutionPlan> { + Arc::new(MockSourceExec::new(make_schema(), num_rows)) + } + + #[test] + fn test_union_provider_sums_rows() -> Result<()> { + let union = UnionExec::try_new(vec![make_source(300), make_source(700)])?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(UnionStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(union.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Exact(1000)); + Ok(()) + } + + #[test] + fn test_union_provider_three_inputs() -> Result<()> { + let union = UnionExec::try_new(vec![ + make_source(100), + make_source(200), + make_source(300), + ])?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(UnionStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(union.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Exact(600)); + Ok(()) + } + + #[test] + fn test_union_provider_absent_propagates() -> Result<()> { + // One input with unknown row count -> result must be Absent, not Inexact(300) + let union = UnionExec::try_new(vec![ + make_source(300), + make_source_with_precision(Precision::Absent), + ])?; + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(UnionStatisticsProvider), + Arc::new(DefaultStatisticsProvider), + ]); + let stats = registry.compute(union.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Absent); + Ok(()) + } + + // ========================================================================= + // ClosureStatisticsProvider tests + // ========================================================================= + + #[test] + fn test_closure_provider_basic() -> Result<()> { + // Override all FilterExec stats with a fixed row count + let provider = ClosureStatisticsProvider::new(|plan, _child_stats| { + if plan.downcast_ref::<FilterExec>().is_some() { + Ok(StatisticsResult::Computed(ExtendedStatistics::from( + Statistics { + num_rows: Precision::Inexact(42), + total_byte_size: Precision::Absent, + column_statistics: vec![], + }, + ))) + } else { + Ok(StatisticsResult::Delegate) + } + }); + + let registry = StatisticsRegistry::with_providers(vec![ + Arc::new(provider), + Arc::new(DefaultStatisticsProvider), + ]); + + let source = make_source(1000); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(lit(true), source)?); + let stats = registry.compute(filter.as_ref())?; + assert_eq!(stats.base.num_rows, Precision::Inexact(42)); + Ok(()) + } + + #[test] + fn test_closure_provider_distinguishes_nodes_by_child_stats() -> Result<()> { + // Two FilterExec nodes with different input sizes. + // The closure uses the child row count as a proxy to distinguish them, + // which mirrors the cardinality feedback use case where you match a + // runtime-observed count to the right node in the plan tree. + let provider = ClosureStatisticsProvider::new(|plan, child_stats| { + if plan.downcast_ref::<FilterExec>().is_none() { + return Ok(StatisticsResult::Delegate); + } + match child_stats[0].base.num_rows.get_value().copied() { + Some(500) => Ok(StatisticsResult::Computed(ExtendedStatistics::from( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Absent, + column_statistics: vec![], + }, + ))), + Some(200) => Ok(StatisticsResult::Computed(ExtendedStatistics::from( + Statistics { + num_rows: Precision::Inexact(50), + total_byte_size: Precision::Absent, + column_statistics: vec![], + }, + ))), + _ => Ok(StatisticsResult::Delegate), + } + }); + + let registry = StatisticsRegistry::with_providers(vec![Arc::new(provider)]); + + let filter_a: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(lit(true), make_source(500))?); + let filter_b: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(lit(true), make_source(200))?); + + let stats_a = registry.compute(filter_a.as_ref())?; + let stats_b = registry.compute(filter_b.as_ref())?; + + assert_eq!(stats_a.base.num_rows, Precision::Inexact(100)); + assert_eq!(stats_b.base.num_rows, Precision::Inexact(50)); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/ordering.rs b/native/vendor/datafusion-physical-plan/src/ordering.rs new file mode 100644 index 0000000..8b596b9 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/ordering.rs
@@ -0,0 +1,54 @@ +// 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. + +/// Specifies how the input to an aggregation or window operator is ordered +/// relative to their `GROUP BY` or `PARTITION BY` expressions. +/// +/// For example, if the existing ordering is `[a ASC, b ASC, c ASC]` +/// +/// ## Window Functions +/// - A `PARTITION BY b` clause can use `Linear` mode. +/// - A `PARTITION BY a, c` or a `PARTITION BY c, a` can use +/// `PartiallySorted([0])` or `PartiallySorted([1])` modes, respectively. +/// (The vector stores the index of `a` in the respective PARTITION BY expression.) +/// - A `PARTITION BY a, b` or a `PARTITION BY b, a` can use `Sorted` mode. +/// +/// ## Aggregations +/// - A `GROUP BY b` clause can use `Linear` mode, as the only one permutation `[b]` +/// cannot satisfy the existing ordering. +/// - A `GROUP BY a, c` or a `GROUP BY c, a` can use +/// `PartiallySorted([0])` or `PartiallySorted([1])` modes, respectively, as +/// the permutation `[a]` satisfies the existing ordering. +/// (The vector stores the index of `a` in the respective PARTITION BY expression.) +/// - A `GROUP BY a, b` or a `GROUP BY b, a` can use `Sorted` mode, as the +/// full permutation `[a, b]` satisfies the existing ordering. +/// +/// Note these are the same examples as above, but with `GROUP BY` instead of +/// `PARTITION BY` to make the examples easier to read. +#[derive(Debug, Clone, PartialEq)] +pub enum InputOrderMode { + /// There is no partial permutation of the expressions satisfying the + /// existing ordering. + Linear, + /// There is a partial permutation of the expressions satisfying the + /// existing ordering. Indices describing the longest partial permutation + /// are stored in the vector. + PartiallySorted(Vec<usize>), + /// There is a (full) permutation of the expressions satisfying the + /// existing ordering. + Sorted, +}
diff --git a/native/vendor/datafusion-physical-plan/src/placeholder_row.rs b/native/vendor/datafusion-physical-plan/src/placeholder_row.rs new file mode 100644 index 0000000..67c063b --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/placeholder_row.rs
@@ -0,0 +1,331 @@ +// 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. + +//! EmptyRelation produce_one_row=true execution plan + +use std::sync::Arc; + +use crate::coop::cooperative; +use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; +use crate::memory::MemoryStream; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, + common, +}; + +use arrow::array::{ArrayRef, NullArray, RecordBatch, RecordBatchOptions}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_or_internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::PhysicalExpr; + +use crate::statistics::StatisticsArgs; +use log::trace; + +/// Execution plan for empty relation with produce_one_row=true +#[derive(Debug, Clone)] +pub struct PlaceholderRowExec { + /// The schema for the produced row + schema: SchemaRef, + /// Number of partitions + partitions: usize, + cache: Arc<PlanProperties>, +} + +impl PlaceholderRowExec { + /// Create a new PlaceholderRowExec + pub fn new(schema: SchemaRef) -> Self { + let partitions = 1; + let cache = Self::compute_properties(Arc::clone(&schema), partitions); + PlaceholderRowExec { + schema, + partitions, + cache: Arc::new(cache), + } + } + + /// Create a new PlaceholderRowExecPlaceholderRowExec with specified partition number + pub fn with_partitions(mut self, partitions: usize) -> Self { + self.partitions = partitions; + // Update output partitioning when updating partitions: + let output_partitioning = Self::output_partitioning_helper(self.partitions); + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + self + } + + fn data(&self) -> Result<Vec<RecordBatch>> { + Ok({ + let n_field = self.schema.fields.len(); + vec![RecordBatch::try_new_with_options( + Arc::new(Schema::new( + (0..n_field) + .map(|i| { + Field::new(format!("placeholder_{i}"), DataType::Null, true) + }) + .collect::<Fields>(), + )), + (0..n_field) + .map(|_i| { + let ret: ArrayRef = Arc::new(NullArray::new(1)); + ret + }) + .collect(), + // Even if column number is empty we can generate single row. + &RecordBatchOptions::new().with_row_count(Some(1)), + )?] + }) + } + + fn output_partitioning_helper(n_partitions: usize) -> Partitioning { + Partitioning::UnknownPartitioning(n_partitions) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef, n_partitions: usize) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Self::output_partitioning_helper(n_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + ) + .with_scheduling_type(SchedulingType::Cooperative) + } +} + +impl DisplayAs for PlaceholderRowExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "PlaceholderRowExec") + } + + DisplayFormatType::TreeRender => Ok(()), + } + } +} + +impl ExecutionPlan for PlaceholderRowExec { + fn name(&self) -> &'static str { + "PlaceholderRowExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start PlaceholderRowExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + assert_or_internal_err!( + partition < self.partitions, + "PlaceholderRowExec invalid partition {partition} (expected less than {})", + self.partitions + ); + + let ms = MemoryStream::try_new(self.data()?, Arc::clone(&self.schema), None)?; + Ok(Box::pin(cooperative(ms))) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let batches = self + .data() + .expect("Create single row placeholder RecordBatch should not fail"); + + let batches = match args.partition() { + Some(_) => vec![batches], + // entire plan + None => vec![batches; self.partitions], + }; + + Ok(Arc::new(common::compute_record_batch_statistics( + &batches, + &self.schema, + None, + ))) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( + protobuf::PlaceholderRowExecNode { + schema: Some(schema), + partitions: self + .properties() + .output_partitioning() + .partition_count() as u32, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl PlaceholderRowExec { + /// Reconstruct a [`PlaceholderRowExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let placeholder = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow, + "PlaceholderRowExec", + ); + let schema = placeholder.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "PlaceholderRowExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(Schema::try_from(schema)?); + // A zero (absent) partition count comes from a plan encoded before the + // field existed, which always meant a single partition. + let partitions = placeholder.partitions.max(1) as usize; + Ok(Arc::new( + PlaceholderRowExec::new(schema).with_partitions(partitions), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{execution_plan::replace_children_if_necessary, test}; + + #[test] + fn replace_children() -> Result<()> { + let schema = test::aggr_test_schema(); + + let placeholder = Arc::new(PlaceholderRowExec::new(schema)); + + let placeholder_2 = replace_children_if_necessary( + Arc::clone(&placeholder) as Arc<dyn ExecutionPlan>, + vec![], + )?; + assert_eq!(placeholder.schema(), placeholder_2.schema()); + + let too_many_kids = vec![placeholder_2]; + assert!( + replace_children_if_necessary(placeholder, too_many_kids).is_err(), + "expected error when providing list of kids" + ); + Ok(()) + } + + #[tokio::test] + async fn invalid_execute() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = test::aggr_test_schema(); + let placeholder = PlaceholderRowExec::new(schema); + + // Ask for the wrong partition + assert!(placeholder.execute(1, Arc::clone(&task_ctx)).is_err()); + assert!(placeholder.execute(20, task_ctx).is_err()); + Ok(()) + } + + #[tokio::test] + async fn produce_one_row() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = test::aggr_test_schema(); + let placeholder = PlaceholderRowExec::new(schema); + + let iter = placeholder.execute(0, task_ctx)?; + let batches = common::collect(iter).await?; + + // Should have one item + assert_eq!(batches.len(), 1); + + Ok(()) + } + + #[tokio::test] + async fn produce_one_row_multiple_partition() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = test::aggr_test_schema(); + let partitions = 3; + let placeholder = PlaceholderRowExec::new(schema).with_partitions(partitions); + + for n in 0..partitions { + let iter = placeholder.execute(n, Arc::clone(&task_ctx))?; + let batches = common::collect(iter).await?; + + // Should have one item + assert_eq!(batches.len(), 1); + } + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/projection.rs b/native/vendor/datafusion-physical-plan/src/projection.rs new file mode 100644 index 0000000..1e672bb --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/projection.rs
@@ -0,0 +1,2523 @@ +// 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. + +//! Defines the projection execution plan. A projection determines which columns or expressions +//! are returned from a query. The SQL statement `SELECT a, b, a+b FROM t1` is an example +//! of a projection on table `t1` where the expressions `a`, `b`, and `a+b` are the +//! projection expressions. `SELECT` without `FROM` will only evaluate expressions. + +use super::expressions::Column; +use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use super::{ + DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, SortOrderPushdownResult, Statistics, +}; +use crate::column_rewriter::PhysicalColumnRewriter; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, FilterRemapper, PushedDownPredicate, +}; +use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, PhysicalExpr, + ReplaceChildrenOptions, validate_child_count, +}; +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; +use datafusion_common::{DataFusionError, JoinSide, Result, internal_err, plan_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::ExpressionPlacement; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr::projection::Projector; +use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; +use datafusion_physical_expr_common::sort_expr::{ + LexOrdering, LexRequirement, PhysicalSortExpr, +}; +// Re-exported from datafusion-physical-expr for backwards compatibility +// We recommend updating your imports to use datafusion-physical-expr directly +pub use datafusion_physical_expr::projection::{ + ProjectionExpr, ProjectionExprs, update_expr, +}; + +use futures::stream::{Stream, StreamExt}; +use log::trace; + +/// [`ExecutionPlan`] for a projection +/// +/// Computes a set of scalar value expressions for each input row, producing one +/// output row for each input row. +#[derive(Debug, Clone)] +pub struct ProjectionExec { + /// A projector specialized to apply the projection to the input schema from the child node + /// and produce [`RecordBatch`]es with the output schema of this node. + projector: Projector, + /// The input plan + input: Arc<dyn ExecutionPlan>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// Whether the output metadata differs from the metadata derived from the + /// projection expressions and input schema. + overrides_metadata: bool, +} + +impl ProjectionExec { + /// Create a projection on an input + /// + /// # Example: + /// Create a `ProjectionExec` to crate `SELECT a, a+b AS sum_ab FROM t1`: + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow_schema::{Schema, Field, DataType}; + /// # use datafusion_expr::Operator; + /// # use datafusion_physical_plan::ExecutionPlan; + /// # use datafusion_physical_expr::expressions::{col, binary}; + /// # use datafusion_physical_plan::empty::EmptyExec; + /// # use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; + /// # fn schema() -> Arc<Schema> { + /// # Arc::new(Schema::new(vec![ + /// # Field::new("a", DataType::Int32, false), + /// # Field::new("b", DataType::Int32, false), + /// # ])) + /// # } + /// # + /// # fn input() -> Arc<dyn ExecutionPlan> { + /// # Arc::new(EmptyExec::new(schema())) + /// # } + /// # + /// # fn main() { + /// let schema = schema(); + /// // Create PhysicalExprs + /// let a = col("a", &schema).unwrap(); + /// let b = col("b", &schema).unwrap(); + /// let a_plus_b = binary(Arc::clone(&a), Operator::Plus, b, &schema).unwrap(); + /// // create ProjectionExec + /// let proj = ProjectionExec::try_new( + /// [ + /// ProjectionExpr { + /// // expr a produces the column named "a" + /// expr: a, + /// alias: "a".to_string(), + /// }, + /// ProjectionExpr { + /// // expr: a + b produces the column named "sum_ab" + /// expr: a_plus_b, + /// alias: "sum_ab".to_string(), + /// }, + /// ], + /// input(), + /// ) + /// .unwrap(); + /// # } + /// ``` + pub fn try_new<I, E>(expr: I, input: Arc<dyn ExecutionPlan>) -> Result<Self> + where + I: IntoIterator<Item = E>, + E: Into<ProjectionExpr>, + { + let input_schema = input.schema(); + let expr_arc = expr.into_iter().map(Into::into).collect::<Arc<_>>(); + let projection = ProjectionExprs::from_expressions(expr_arc); + let projector = projection.make_projector(&input_schema)?; + Self::try_from_projector(projector, input, false) + } + + /// Create a projection using field and schema metadata from + /// `projected_schema`. + /// + /// Field names, data types, and nullability are still derived from the physical + /// projection expressions and the input plan; only field and schema metadata are + /// taken from `projected_schema`. + /// + /// # Errors + /// + /// Returns an error if the projection cannot be applied to the input plan, or if + /// `projected_schema` has a different number of fields than the projection. + pub fn try_new_with_schema_metadata<I, E>( + expr: I, + input: Arc<dyn ExecutionPlan>, + projected_schema: &Schema, + ) -> Result<Self> + where + I: IntoIterator<Item = E>, + E: Into<ProjectionExpr>, + { + let input_schema = input.schema(); + let expr_arc = expr.into_iter().map(Into::into).collect::<Arc<_>>(); + let projection = ProjectionExprs::from_expressions(expr_arc); + let projector = projection + .make_projector_with_schema_metadata(&input_schema, projected_schema)?; + let overrides_metadata = + Self::compute_overrides_metadata(&projector, &input_schema)?; + Self::try_from_projector(projector, input, overrides_metadata) + } + + fn try_from_projector( + projector: Projector, + input: Arc<dyn ExecutionPlan>, + overrides_metadata: bool, + ) -> Result<Self> { + // Construct a map from the input expressions to the output expression of the Projection + let projection_mapping = + projector.projection().projection_mapping(&input.schema())?; + let cache = Self::compute_properties( + &input, + &projection_mapping, + Arc::clone(projector.output_schema()), + )?; + Ok(Self { + projector, + input, + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + overrides_metadata, + }) + } + + /// The projection expressions stored as tuples of (expression, output column name) + pub fn expr(&self) -> &[ProjectionExpr] { + self.projector.projection().as_ref() + } + + /// The projection expressions as a [`ProjectionExprs`]. + pub fn projection_expr(&self) -> &ProjectionExprs { + self.projector.projection() + } + + /// The input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + projection_mapping: &ProjectionMapping, + schema: SchemaRef, + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let input_eq_properties = input.equivalence_properties(); + let eq_properties = input_eq_properties.project(projection_mapping, schema); + // Calculate output partitioning, which needs to respect aliases: + let output_partitioning = input + .output_partitioning() + .project(projection_mapping, input_eq_properties); + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + input.pipeline_behavior(), + input.boundedness(), + )) + } + + /// Returns whether `projector`'s output metadata differs from the metadata + /// derived from its expressions and `input_schema`. + fn compute_overrides_metadata( + projector: &Projector, + input_schema: &Schema, + ) -> Result<bool> { + let output_schema = projector.output_schema(); + if input_schema.metadata() != output_schema.metadata() { + return Ok(true); + } + for (projection, output_field) in + projector.projection().iter().zip(output_schema.fields()) + { + let derived_field = projection.expr.return_field(input_schema)?; + if derived_field.metadata() != output_field.metadata() { + return Ok(true); + } + } + Ok(false) + } + + /// Returns whether this projection's output metadata differs from the + /// metadata derived when the projection was constructed. + fn overrides_metadata(&self) -> bool { + self.overrides_metadata + } + + /// Collect reverse alias mapping from projection expressions. + /// The result hash map is a map from aliased Column in parent to original expr. + fn collect_reverse_alias( + &self, + ) -> Result<datafusion_common::HashMap<Column, Arc<dyn PhysicalExpr>>> { + let mut alias_map = datafusion_common::HashMap::new(); + for projection in self.projection_expr().iter() { + let (aliased_index, _output_field) = self + .projector + .output_schema() + .column_with_name(&projection.alias) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "Expr {} with alias {} not found in output schema", + projection.expr, projection.alias + )) + })?; + let aliased_col = Column::new(&projection.alias, aliased_index); + alias_map.insert(aliased_col, Arc::clone(&projection.expr)); + } + Ok(alias_map) + } +} + +impl DisplayAs for ProjectionExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let expr: Vec<String> = self + .projector + .projection() + .as_ref() + .iter() + .map(|proj_expr| { + let e = proj_expr.expr.to_string(); + if e != proj_expr.alias { + format!("{e} as {}", proj_expr.alias) + } else { + e + } + }) + .collect(); + + write!(f, "ProjectionExec: expr=[{}]", expr.join(", ")) + } + DisplayFormatType::TreeRender => { + for (i, proj_expr) in self.expr().iter().enumerate() { + let expr_sql = fmt_sql(proj_expr.expr.as_ref()); + if proj_expr.expr.to_string() == proj_expr.alias { + writeln!(f, "expr{i}={expr_sql}")?; + } else { + writeln!(f, "{}={expr_sql}", proj_expr.alias)?; + } + } + + Ok(()) + } + } + } +} + +impl ExecutionPlan for ProjectionExec { + fn name(&self) -> &'static str { + "ProjectionExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn maintains_input_order(&self) -> Vec<bool> { + // Tell optimizer this operator doesn't reorder its input + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + let all_simple_exprs = + self.projector + .projection() + .as_ref() + .iter() + .all(|proj_expr| { + !matches!( + proj_expr.expr.placement(), + ExpressionPlacement::KeepInPlace + ) + }); + // If expressions are all either column_expr or Literal (or other cheap expressions), + // then all computations in this projection are reorder or rename, + // and projection would not benefit from the repartition, benefits_from_input_partitioning will return false. + vec![!all_simple_exprs] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + crate::apply_expression_roots(self.projector.projection().as_ref().iter(), f) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let input = children.swap_remove(0); + let projector = self.projector.clone(); + let overrides_metadata = ProjectionExec::compute_overrides_metadata( + &projector, + input.schema().as_ref(), + )?; + ProjectionExec::try_from_projector(projector, input, overrides_metadata) + .map(|p| Arc::new(p) as _) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start ProjectionExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + let projector = self.projector.with_metrics(&self.metrics, partition); + Ok(Box::pin(ProjectionStream::new( + projector, + self.input.execute(partition, context)?, + BaselineMetrics::new(&self.metrics, partition), + )?)) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let input_stats = input_stats[0].as_ref().clone(); + let output_schema = self.schema(); + Ok(Arc::new( + self.projector + .projection() + .project_statistics(input_stats, &output_schema)?, + )) + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + match try_collapse_projection_chain(projection)? { + Some(plan) => Ok(Some(plan)), + None => Ok(Some(Arc::new(projection.clone()))), + } + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + // expand alias column to original expr in parent filters + let invert_alias_map = self.collect_reverse_alias()?; + let output_schema = self.schema(); + let remapper = FilterRemapper::new(output_schema); + let mut child_parent_filters = Vec::with_capacity(parent_filters.len()); + + for filter in parent_filters { + // Check that column exists in child, then reassign column indices to match child schema + if let Some(reassigned) = remapper.try_remap(&filter)? { + // rewrite filter expression using invert alias map + let mut rewriter = PhysicalColumnRewriter::new(&invert_alias_map); + let rewritten = reassigned.rewrite(&mut rewriter)?.data; + child_parent_filters.push(PushedDownPredicate::supported(rewritten)); + } else { + child_parent_filters.push(PushedDownPredicate::unsupported(filter)); + } + } + + Ok(FilterDescription::new().with_child(ChildFilterDescription { + parent_filters: child_parent_filters, + self_filters: vec![], + })) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + Ok(FilterPushdownPropagation::if_all(child_pushdown_result)) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + let child = self.input(); + let mut child_order = Vec::new(); + + // Check and transform sort expressions + for sort_expr in order { + // Recursively transform the expression + let mut can_pushdown = true; + let transformed = Arc::clone(&sort_expr.expr).transform(|expr| { + if let Some(col) = expr.downcast_ref::<Column>() { + // Check if column index is valid. + // This should always be true but fail gracefully if it's not. + if col.index() >= self.expr().len() { + can_pushdown = false; + return Ok(Transformed::no(expr)); + } + + let proj_expr = &self.expr()[col.index()]; + + // Check if projection expression is a simple column + // We cannot push down order by clauses that depend on + // projected computations as they would have nothing to reference. + if let Some(child_col) = proj_expr.expr.downcast_ref::<Column>() { + // Replace with the child column + Ok(Transformed::yes(Arc::new(child_col.clone()) as _)) + } else { + // Projection involves computation, cannot push down + can_pushdown = false; + Ok(Transformed::no(expr)) + } + } else { + Ok(Transformed::no(expr)) + } + })?; + + if !can_pushdown { + return Ok(SortOrderPushdownResult::Unsupported); + } + + child_order.push(PhysicalSortExpr { + expr: transformed.data, + options: sort_expr.options, + }); + } + + // Recursively push down to child node + match child.try_pushdown_sort(&child_order)? { + SortOrderPushdownResult::Exact { inner } => { + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; + Ok(SortOrderPushdownResult::Exact { inner: new_exec }) + } + SortOrderPushdownResult::Inexact { inner } => { + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; + Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) + } + SortOrderPushdownResult::Unsupported => { + Ok(SortOrderPushdownResult::Unsupported) + } + } + } + + fn with_preserve_order( + &self, + preserve_order: bool, + ) -> Option<Arc<dyn ExecutionPlan>> { + self.input + .with_preserve_order(preserve_order) + .and_then(|new_input| { + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) + .ok() + }) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expressions(self.expr().iter().map(|p| &p.expr))?; + let expr_name = self.expr().iter().map(|p| p.alias.clone()).collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Projection(Box::new( + protobuf::ProjectionExecNode { + input: Some(Box::new(input)), + expr, + expr_name, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ProjectionExec { + /// Reconstruct a [`ProjectionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. Child plans and expressions are decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let projection = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Projection, + "ProjectionExec", + ); + let input = ctx.decode_required_child( + projection.input.as_deref(), + "ProjectionExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = projection + .expr + .iter() + .zip(projection.expr_name.iter()) + .map(|(expr, name)| { + Ok(ProjectionExpr { + expr: ctx.decode_expr(expr, input_schema.as_ref())?, + alias: name.to_string(), + }) + }) + .collect::<Result<Vec<ProjectionExpr>>>()?; + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) + } +} + +impl ProjectionStream { + /// Create a new projection stream + fn new( + projector: Projector, + input: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, + ) -> Result<Self> { + Ok(Self { + projector, + input, + baseline_metrics, + }) + } + + fn batch_project(&self, batch: &RecordBatch) -> Result<RecordBatch> { + // Records time on drop + let _timer = self.baseline_metrics.elapsed_compute().timer(); + self.projector.project_batch(batch) + } +} + +/// Projection iterator +struct ProjectionStream { + projector: Projector, + input: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, +} + +impl Stream for ProjectionStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let poll = self.input.poll_next_unpin(cx).map(|x| match x { + Some(Ok(batch)) => Some(self.batch_project(&batch)), + other => other, + }); + + self.baseline_metrics.record_poll(poll) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + // Same number of record batches + self.input.size_hint() + } +} + +impl RecordBatchStream for ProjectionStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(self.projector.output_schema()) + } +} + +/// Trait for execution plans that can embed a projection, avoiding a separate +/// [`ProjectionExec`] wrapper. +/// +/// # Empty projections +/// +/// `Some(vec![])` is a valid projection that produces zero output columns while +/// preserving the correct row count. Implementors must ensure that runtime batch +/// construction still returns batches with the right number of rows even when no +/// columns are selected (e.g. for `SELECT count(1) … JOIN …`). +pub trait EmbeddedProjection: ExecutionPlan + Sized { + fn with_projection(&self, projection: Option<Vec<usize>>) -> Result<Self>; +} + +/// Some projection can't be pushed down left input or right input of hash join because filter or on need may need some columns that won't be used in later. +/// By embed those projection to hash join, we can reduce the cost of build_batch_from_indices in hash join (build_batch_from_indices need to can compute::take() for each column) and avoid unnecessary output creation. +pub fn try_embed_projection<Exec: EmbeddedProjection + 'static>( + projection: &ProjectionExec, + execution_plan: &Exec, +) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection has no expressions at all (e.g., ProjectionExec: expr=[]), + // embed an empty projection into the execution plan so it outputs zero columns. + // This avoids allocating throwaway null arrays for build-side columns + // when no output columns are actually needed (e.g., count(1) over a right join). + if projection.expr().is_empty() { + let new_execution_plan = Arc::new(execution_plan.with_projection(Some(vec![]))?); + return Ok(Some(new_execution_plan)); + } + + // Collect all column indices from the given projection expressions. + let projection_index = collect_column_indices(projection.expr()); + + if projection_index.is_empty() { + return Ok(None); + }; + + let columns_reduced = projection_index.len() < execution_plan.schema().fields().len(); + + let new_execution_plan = + Arc::new(execution_plan.with_projection(Some(projection_index.to_vec()))?); + + // Build projection expressions for update_expr. Zip the projection_index with the new_execution_plan output schema fields. + let embed_project_exprs = projection_index + .iter() + .zip(new_execution_plan.schema().fields()) + .map(|(index, field)| ProjectionExpr { + expr: Arc::new(Column::new(field.name(), *index)) as Arc<dyn PhysicalExpr>, + alias: field.name().to_owned(), + }) + .collect::<Vec<_>>(); + + let mut new_projection_exprs = Vec::with_capacity(projection.expr().len()); + + for proj_expr in projection.expr() { + // update column index for projection expression since the input schema has been changed. + let Some(expr) = + update_expr(&proj_expr.expr, embed_project_exprs.as_slice(), false)? + else { + return Ok(None); + }; + new_projection_exprs.push(ProjectionExpr { + expr, + alias: proj_expr.alias.clone(), + }); + } + // Old projection may contain some alias or expression such as `a + 1` and `CAST('true' AS BOOLEAN)`, but our projection_exprs in hash join just contain column, so we need to create the new projection to keep the original projection. + let new_projection = Arc::new(ProjectionExec::try_new( + new_projection_exprs, + Arc::clone(&new_execution_plan) as _, + )?); + if is_projection_removable(&new_projection) { + // Residual is identity — embedding fully absorbed the projection. + Ok(Some(new_execution_plan)) + } else if columns_reduced { + // Embedding reduced columns even though a residual is still needed + // for renames or expressions — worth keeping. + Ok(Some(new_projection)) + } else { + // No columns eliminated and residual still needed — embedding just + // adds an unnecessary column reorder inside the operator. + Ok(None) + } +} + +pub struct JoinData { + pub projected_left_child: ProjectionExec, + pub projected_right_child: ProjectionExec, + pub join_filter: Option<JoinFilter>, + pub join_on: JoinOn, +} + +#[deprecated( + since = "55.0.0", + note = "Use try_pushdown_through_join_with_column_indices instead" +)] +pub fn try_pushdown_through_join( + projection: &ProjectionExec, + join_left: &Arc<dyn ExecutionPlan>, + join_right: &Arc<dyn ExecutionPlan>, + join_on: JoinOnRef, + schema: &SchemaRef, + filter: Option<&JoinFilter>, +) -> Result<Option<JoinData>> { + let left_field_count = join_left.schema().fields().len(); + let column_indices = schema + .fields() + .iter() + .enumerate() + .map(|(index, _)| { + if index < left_field_count { + ColumnIndex { + index, + side: JoinSide::Left, + } + } else { + ColumnIndex { + index: index - left_field_count, + side: JoinSide::Right, + } + } + }) + .collect::<Vec<_>>(); + + try_pushdown_through_join_with_column_indices( + projection, + join_left, + join_right, + join_on, + schema, + filter, + &column_indices, + ) +} + +/// Attempts to move a projection below a join by mapping each join output +/// column to the child column that produced it. +/// +/// `schema` is the complete output schema of the join, not either child's +/// schema. `column_indices` must contain one entry for each field in `schema`. +/// Each [`JoinSide::Left`] or [`JoinSide::Right`] entry identifies the source +/// child and uses an index relative to that child's schema. +/// +/// [`JoinSide::None`] identifies a column produced by the join itself, such as +/// a mark column. If `projection` references such a column, this function +/// returns `Ok(None)` because neither child can produce it. +/// +/// Returns `Ok(None)` when the projection cannot be pushed down safely. +/// +/// # Errors +/// +/// Returns an error if `column_indices` does not match `schema` or contains an +/// index outside the corresponding child schema. +pub fn try_pushdown_through_join_with_column_indices( + projection: &ProjectionExec, + join_left: &Arc<dyn ExecutionPlan>, + join_right: &Arc<dyn ExecutionPlan>, + join_on: JoinOnRef, + schema: &SchemaRef, + filter: Option<&JoinFilter>, + column_indices: &[ColumnIndex], +) -> Result<Option<JoinData>> { + if column_indices.len() != schema.fields().len() { + return plan_err!( + "Column index mapping has {} entries but join schema has {} fields", + column_indices.len(), + schema.fields().len() + ); + } + // Validate each output-to-child mapping before using it to rewrite the + // projection. Synthetic outputs have no child index to validate. + for (output_index, column_index) in column_indices.iter().enumerate() { + let (side, child_field_count) = match column_index.side { + JoinSide::Left => ("left", join_left.schema().fields().len()), + JoinSide::Right => ("right", join_right.schema().fields().len()), + JoinSide::None => continue, + }; + if column_index.index >= child_field_count { + return plan_err!( + "Join output column {output_index} maps to {side} child column {}, but the child has {child_field_count} fields", + column_index.index + ); + } + } + + // Convert projected expressions to columns. We can not proceed if this is not possible. + let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) else { + return Ok(None); + }; + + if projection_as_columns.len() >= schema.fields().len() { + return Ok(None); + } + let mut left_proj: Vec<(Column, String)> = Vec::new(); + let mut right_proj: Vec<(Column, String)> = Vec::new(); + let mut seen_right = false; + for (col, alias) in &projection_as_columns { + let Some(origin) = column_indices.get(col.index()) else { + return plan_err!( + "Projection column {} is outside the {}-entry column index mapping", + col.index(), + column_indices.len() + ); + }; + match origin.side { + // Keep the "left block before right block" contiguity the current + // pushdown supports; a left column after a right one is "mixed". + JoinSide::Left => { + if seen_right { + return Ok(None); + } + left_proj.push((Column::new(col.name(), origin.index), alias.clone())); + } + JoinSide::Right => { + seen_right = true; + right_proj.push((Column::new(col.name(), origin.index), alias.clone())); + } + // Synthetic column (e.g. mark): belongs to neither child. + // Phase 2 declines; Phase 3 keeps it at the join output instead. + JoinSide::None => return Ok(None), + } + } + + // Parity: neither side fully dropped. + if left_proj.is_empty() || right_proj.is_empty() { + return Ok(None); + } + + // `left_proj` / `right_proj` carry *child* indices (from `column_indices`), + // so the shared `update_join_*` helpers must use a 0 column-index offset for + // both sides (the offset bridges child -> join-output index, which is the + // identity here). + let new_filter = if let Some(filter) = filter { + match update_join_filter(&left_proj, &right_proj, filter, 0) { + Some(updated) => Some(updated), + None => return Ok(None), + } + } else { + None + }; + + let Some(new_on) = update_join_on(&left_proj, &right_proj, join_on, 0) else { + return Ok(None); + }; + + let (new_left, new_right) = + new_join_children_from_groups(&left_proj, &right_proj, join_left, join_right)?; + + Ok(Some(JoinData { + projected_left_child: new_left, + projected_right_child: new_right, + join_filter: new_filter, + join_on: new_on, + })) +} + +/// This function checks if `plan` is a [`ProjectionExec`], and inspects its +/// input(s) to test whether it can push `plan` under its input(s). This function +/// will operate on the entire tree and may ultimately remove `plan` entirely +/// by leveraging source providers with built-in projection capabilities. +pub fn remove_unnecessary_projections( + plan: Arc<dyn ExecutionPlan>, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let maybe_modified = if let Some(projection) = plan.downcast_ref::<ProjectionExec>() { + // If the projection does not cause any change on the input, we can + // safely remove it: + if is_projection_removable(projection) { + return Ok(Transformed::yes(Arc::clone(projection.input()))); + } + // Swapping a projection with observable metadata can change query results + // by changing the metadata visible to its child expressions. + if projection.overrides_metadata() { + return Ok(Transformed::no(plan)); + } + // Otherwise, check if we can push it under its child(ren): + projection + .input() + .try_swapping_with_projection(projection)? + } else { + return Ok(Transformed::no(plan)); + }; + Ok(maybe_modified.map_or_else(|| Transformed::no(plan), Transformed::yes)) +} + +/// Compare the inputs and outputs of the projection. All expressions must be +/// columns without alias, and projection does not change the order of fields. +/// The input and output schemas must also match exactly to preserve metadata. +/// For example, if the input schema is `a, b`, `SELECT a, b` is removable, +/// but `SELECT b, a` and `SELECT a+1, b` and `SELECT a AS c, b` are not. +fn is_projection_removable(projection: &ProjectionExec) -> bool { + let exprs = projection.expr(); + exprs.iter().enumerate().all(|(idx, proj_expr)| { + let Some(col) = proj_expr.expr.downcast_ref::<Column>() else { + return false; + }; + col.name() == proj_expr.alias && col.index() == idx + }) && exprs.len() == projection.input().schema().fields().len() + && projection.schema() == projection.input().schema() +} + +/// Given the expression set of a projection, checks if the projection causes +/// any renaming or constructs a non-`Column` physical expression. +pub fn all_alias_free_columns(exprs: &[ProjectionExpr]) -> bool { + exprs.iter().all(|proj_expr| { + proj_expr + .expr + .downcast_ref::<Column>() + .map(|column| column.name() == proj_expr.alias) + .unwrap_or(false) + }) +} + +/// Updates a source provider's projected columns according to the given +/// projection operator's expressions. To use this function safely, one must +/// ensure that all expressions are `Column` expressions without aliases. +pub fn new_projections_for_columns( + projection: &[ProjectionExpr], + source: &[usize], +) -> Vec<usize> { + projection + .iter() + .filter_map(|proj_expr| { + proj_expr + .expr + .downcast_ref::<Column>() + .map(|expr| source[expr.index()]) + }) + .collect() +} + +/// Creates a new [`ProjectionExec`] instance with the given child plan and +/// projected expressions, preserving the original output metadata. +pub fn make_with_child( + projection: &ProjectionExec, + child: &Arc<dyn ExecutionPlan>, +) -> Result<Arc<dyn ExecutionPlan>> { + ProjectionExec::try_new_with_schema_metadata( + projection.expr().to_vec(), + Arc::clone(child), + projection.schema().as_ref(), + ) + .map(|e| Arc::new(e) as _) +} + +/// Returns `true` if all the expressions in the argument are `Column`s. +pub fn all_columns(exprs: &[ProjectionExpr]) -> bool { + exprs.iter().all(|proj_expr| proj_expr.expr.is::<Column>()) +} + +/// Updates the given lexicographic ordering according to given projected +/// expressions using the [`update_expr`] function. +pub fn update_ordering( + ordering: LexOrdering, + projected_exprs: &[ProjectionExpr], +) -> Result<Option<LexOrdering>> { + let mut updated_exprs = vec![]; + for mut sort_expr in ordering.into_iter() { + let Some(updated_expr) = update_expr(&sort_expr.expr, projected_exprs, false)? + else { + return Ok(None); + }; + sort_expr.expr = updated_expr; + updated_exprs.push(sort_expr); + } + Ok(LexOrdering::new(updated_exprs)) +} + +/// Updates the given lexicographic requirement according to given projected +/// expressions using the [`update_expr`] function. +pub fn update_ordering_requirement( + reqs: LexRequirement, + projected_exprs: &[ProjectionExpr], +) -> Result<Option<LexRequirement>> { + let mut updated_exprs = vec![]; + for mut sort_expr in reqs.into_iter() { + let Some(updated_expr) = update_expr(&sort_expr.expr, projected_exprs, false)? + else { + return Ok(None); + }; + sort_expr.expr = updated_expr; + updated_exprs.push(sort_expr); + } + Ok(LexRequirement::new(updated_exprs)) +} + +/// Downcasts all the expressions in `exprs` to `Column`s. If any of the given +/// expressions is not a `Column`, returns `None`. +pub fn physical_to_column_exprs( + exprs: &[ProjectionExpr], +) -> Option<Vec<(Column, String)>> { + exprs + .iter() + .map(|proj_expr| { + proj_expr + .expr + .downcast_ref::<Column>() + .map(|col| (col.clone(), proj_expr.alias.clone())) + }) + .collect() +} + +/// If pushing down the projection over this join's children seems possible, +/// this function constructs the new [`ProjectionExec`]s that will come on top +/// of the original children of the join. +pub fn new_join_children( + projection_as_columns: &[(Column, String)], + far_right_left_col_ind: i32, + far_left_right_col_ind: i32, + left_child: &Arc<dyn ExecutionPlan>, + right_child: &Arc<dyn ExecutionPlan>, +) -> Result<(ProjectionExec, ProjectionExec)> { + let new_left = ProjectionExec::try_new( + projection_as_columns[0..=far_right_left_col_ind as _] + .iter() + .map(|(col, alias)| ProjectionExpr { + expr: Arc::new(Column::new(col.name(), col.index())) as _, + alias: alias.clone(), + }), + Arc::clone(left_child), + )?; + let left_size = left_child.schema().fields().len() as i32; + let new_right = ProjectionExec::try_new( + projection_as_columns[far_left_right_col_ind as _..] + .iter() + .map(|(col, alias)| { + ProjectionExpr { + expr: Arc::new(Column::new( + col.name(), + // Align projected expressions coming from the right + // table with the new right child projection: + (col.index() as i32 - left_size) as _, + )) as _, + alias: alias.clone(), + } + }), + Arc::clone(right_child), + )?; + + Ok((new_left, new_right)) +} + +/// Build the projected left and right children from side-grouped projection +/// columns whose indices are already *child*-relative (e.g. derived from a +/// join's `ColumnIndex`). Unlike [`new_join_children`], this does not infer +/// child ownership from output position, so it is safe for join schemas whose +/// output is not a plain `left ++ right` (used by the schema-aware +/// `try_pushdown_through_join_with_column_indices`). +fn new_join_children_from_groups( + left_proj: &[(Column, String)], + right_proj: &[(Column, String)], + left_child: &Arc<dyn ExecutionPlan>, + right_child: &Arc<dyn ExecutionPlan>, +) -> Result<(ProjectionExec, ProjectionExec)> { + let build = |cols: &[(Column, String)], child: &Arc<dyn ExecutionPlan>| { + ProjectionExec::try_new( + cols.iter().map(|(col, alias)| ProjectionExpr { + expr: Arc::new(Column::new(col.name(), col.index())) as _, + alias: alias.clone(), + }), + Arc::clone(child), + ) + }; + + Ok(( + build(left_proj, left_child)?, + build(right_proj, right_child)?, + )) +} + +/// Checks three conditions for pushing a projection down through a join: +/// - Projection must narrow the join output schema. +/// - Columns coming from left/right tables must be collected at the left/right +/// sides of the output table. +/// - Left or right table is not lost after the projection. +pub fn join_allows_pushdown( + projection_as_columns: &[(Column, String)], + join_schema: &SchemaRef, + far_right_left_col_ind: i32, + far_left_right_col_ind: i32, +) -> bool { + // Projection must narrow the join output: + projection_as_columns.len() < join_schema.fields().len() + // Are the columns from different tables mixed? + && (far_right_left_col_ind + 1 == far_left_right_col_ind) + // Left or right table is not lost after the projection. + && far_right_left_col_ind >= 0 + && far_left_right_col_ind < projection_as_columns.len() as i32 +} + +/// Returns the last index before encountering a column coming from the right table when traveling +/// through the projection from left to right, and the last index before encountering a column +/// coming from the left table when traveling through the projection from right to left. +/// If there is no column in the projection coming from the left side, it returns (-1, ...), +/// if there is no column in the projection coming from the right side, it returns (..., projection length). +pub fn join_table_borders( + left_table_column_count: usize, + projection_as_columns: &[(Column, String)], +) -> (i32, i32) { + let far_right_left_col_ind = projection_as_columns + .iter() + .enumerate() + .take_while(|(_, (projection_column, _))| { + projection_column.index() < left_table_column_count + }) + .last() + .map(|(index, _)| index as i32) + .unwrap_or(-1); + + let far_left_right_col_ind = projection_as_columns + .iter() + .enumerate() + .rev() + .take_while(|(_, (projection_column, _))| { + projection_column.index() >= left_table_column_count + }) + .last() + .map(|(index, _)| index as i32) + .unwrap_or(projection_as_columns.len() as i32); + + (far_right_left_col_ind, far_left_right_col_ind) +} + +/// Tries to update the equi-join `Column`'s of a join as if the input of +/// the join was replaced by a projection. +pub fn update_join_on( + proj_left_exprs: &[(Column, String)], + proj_right_exprs: &[(Column, String)], + hash_join_on: &[(PhysicalExprRef, PhysicalExprRef)], + left_field_size: usize, +) -> Option<Vec<(PhysicalExprRef, PhysicalExprRef)>> { + let (left_idx, right_idx): (Vec<_>, Vec<_>) = hash_join_on + .iter() + .map(|(left, right)| (left, right)) + .unzip(); + + let new_left = new_columns_for_join_on(&left_idx, proj_left_exprs, 0)?; + let new_right = + new_columns_for_join_on(&right_idx, proj_right_exprs, left_field_size)?; + Some(new_left.into_iter().zip(new_right).collect()) +} + +/// Tries to update the column indices of a [`JoinFilter`] as if the input of +/// the join was replaced by a projection. +pub fn update_join_filter( + projection_left_exprs: &[(Column, String)], + projection_right_exprs: &[(Column, String)], + join_filter: &JoinFilter, + left_field_size: usize, +) -> Option<JoinFilter> { + let mut new_left_indices = new_indices_for_join_filter( + join_filter, + JoinSide::Left, + projection_left_exprs, + 0, + ) + .into_iter(); + let mut new_right_indices = new_indices_for_join_filter( + join_filter, + JoinSide::Right, + projection_right_exprs, + left_field_size, + ) + .into_iter(); + + // Check if all columns match: + (new_right_indices.len() + new_left_indices.len() + == join_filter.column_indices().len()) + .then(|| { + JoinFilter::new( + Arc::clone(join_filter.expression()), + join_filter + .column_indices() + .iter() + .map(|col_idx| ColumnIndex { + index: if col_idx.side == JoinSide::Left { + new_left_indices.next().unwrap() + } else { + new_right_indices.next().unwrap() + }, + side: col_idx.side, + }) + .collect(), + Arc::clone(join_filter.schema()), + ) + }) +} + +/// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns +/// `None` if nothing could be merged. +/// +/// The projection-removal optimizer checks `outer.overrides_metadata()` before +/// reaching this helper. The unified projection also keeps `outer`'s schema, so +/// collapsing cannot lose its output metadata. Inner projections still need the +/// check below because outer expressions may observe their metadata. +fn try_collapse_projection_chain( + outer: &ProjectionExec, +) -> Result<Option<Arc<dyn ExecutionPlan>>> { + let mut current_exprs: Vec<ProjectionExpr> = outer.expr().to_vec(); + let mut current_input: Arc<dyn ExecutionPlan> = Arc::clone(outer.input()); + let mut column_ref_map: HashMap<Column, usize> = HashMap::new(); + let mut collapsed_any = false; + + 'outer: while let Some(inner_proj) = current_input.downcast_ref::<ProjectionExec>() { + if inner_proj.overrides_metadata() { + break; + } + + // Collect the column references usage in the outer projection. + column_ref_map.clear(); + for proj_expr in ¤t_exprs { + proj_expr.expr.apply(|expr| { + if let Some(column) = expr.downcast_ref::<Column>() { + *column_ref_map.entry(column.clone()).or_default() += 1; + } + Ok(TreeNodeRecursion::Continue) + })?; + } + let inner_exprs = inner_proj.expr(); + // Merging these projections is not beneficial, e.g + // If an expression is not trivial (KeepInPlace) and it is referred more than 1, unifies projections will be + // beneficial as caching mechanism for non-trivial computations. + // See discussion in: https://github.com/apache/datafusion/issues/8296 + let blocked = column_ref_map.iter().any(|(column, count)| { + *count > 1 + && !inner_exprs[column.index()] + .expr + .placement() + .should_push_to_leaves() + }); + if blocked { + break; + } + + let mut new_phys: Vec<Arc<dyn PhysicalExpr>> = + Vec::with_capacity(current_exprs.len()); + for proj_expr in ¤t_exprs { + // If there is no match in the input projection, we cannot unify these + // projections. This case will arise if the projection expression contains + // a `PhysicalExpr` variant `update_expr` doesn't support. + let Some(expr) = update_expr(&proj_expr.expr, inner_exprs, true)? else { + break 'outer; + }; + new_phys.push(expr); + } + for (proj_expr, expr) in current_exprs.iter_mut().zip(new_phys) { + proj_expr.expr = expr; + } + current_input = Arc::clone(inner_proj.input()); + collapsed_any = true; + } + + if !collapsed_any { + return Ok(None); + } + + // To unify 3 or more sequential projections: + // Preserve the outer projection's output metadata. + let unified: Arc<dyn ExecutionPlan> = + Arc::new(ProjectionExec::try_new_with_schema_metadata( + current_exprs, + current_input, + outer.schema().as_ref(), + )?); + remove_unnecessary_projections(unified).data().map(Some) +} + +/// Collect all column indices from the given projection expressions. +fn collect_column_indices(exprs: &[ProjectionExpr]) -> Vec<usize> { + // Collect column indices in a deterministic order that preserves the + // projection's column ordering. For simple Column expressions, we use + // the column index directly. For complex expressions, we walk the + // expression tree to collect column references in traversal order. + // This allows the embedded projection to match the desired output + // column order, avoiding a residual ProjectionExec. + let mut seen = std::collections::HashSet::new(); + let mut indices = Vec::new(); + for proj_expr in exprs { + if let Some(col) = proj_expr.expr.downcast_ref::<Column>() { + // Simple column reference: preserve projection order. + if seen.insert(col.index()) { + indices.push(col.index()); + } + } else { + // Complex expression: collect all referenced columns in + // expression tree traversal order (deterministic) to preserve + // the natural ordering of column references. + proj_expr + .expr + .apply(|expr| { + if let Some(col) = expr.downcast_ref::<Column>() + && seen.insert(col.index()) + { + indices.push(col.index()); + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("closure always returns OK"); + } + } + indices +} + +/// This function determines and returns a vector of indices representing the +/// positions of columns in `projection_exprs` that are involved in `join_filter`, +/// and correspond to a particular side (`join_side`) of the join operation. +/// +/// Notes: Column indices in the projection expressions are based on the join schema, +/// whereas the join filter is based on the join child schema. `column_index_offset` +/// represents the offset between them. +fn new_indices_for_join_filter( + join_filter: &JoinFilter, + join_side: JoinSide, + projection_exprs: &[(Column, String)], + column_index_offset: usize, +) -> Vec<usize> { + join_filter + .column_indices() + .iter() + .filter(|col_idx| col_idx.side == join_side) + .filter_map(|col_idx| { + projection_exprs + .iter() + .position(|(col, _)| col_idx.index + column_index_offset == col.index()) + }) + .collect() +} + +/// This function generates a new set of columns to be used in a hash join +/// operation based on a set of equi-join conditions (`hash_join_on`) and a +/// list of projection expressions (`projection_exprs`). +/// +/// Notes: Column indices in the projection expressions are based on the join schema, +/// whereas the join on expressions are based on the join child schema. `column_index_offset` +/// represents the offset between them. +fn new_columns_for_join_on( + hash_join_on: &[&PhysicalExprRef], + projection_exprs: &[(Column, String)], + column_index_offset: usize, +) -> Option<Vec<PhysicalExprRef>> { + let new_columns = hash_join_on + .iter() + .filter_map(|on| { + // Rewrite all columns in `on` + Arc::clone(*on) + .transform(|expr| { + if let Some(column) = expr.downcast_ref::<Column>() { + // Find the column in the projection expressions + let new_column = projection_exprs + .iter() + .enumerate() + .find(|(_, (proj_column, _))| { + column.name() == proj_column.name() + && column.index() + column_index_offset + == proj_column.index() + }) + .map(|(index, (_, alias))| Column::new(alias, index)); + if let Some(new_column) = new_column { + Ok(Transformed::yes(Arc::new(new_column))) + } else { + // If the column is not found in the projection expressions, + // it means that the column is not projected. In this case, + // we cannot push the projection down. + internal_err!( + "Column {:?} not found in projection expressions", + column + ) + } + } else { + Ok(Transformed::no(expr)) + } + }) + .data() + .ok() + }) + .collect::<Vec<_>>(); + (new_columns.len() == hash_join_on.len()).then_some(new_columns) +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::common::collect; + use crate::empty::EmptyExec; + use crate::filter::FilterExec; + + use crate::filter_pushdown::PushedDown; + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::test; + use crate::test::exec::StatisticsExec; + + use arrow::array::StringArray; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::ScalarValue; + use datafusion_common::stats::{ColumnStatistics, Precision, Statistics}; + + use datafusion_expr::{Operator, ScalarUDF}; + use datafusion_functions::core::arrow_metadata::ArrowMetadataFunc; + use datafusion_physical_expr::ScalarFunctionExpr; + use datafusion_physical_expr::expressions::{ + BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, is_null, lit, + }; + + #[test] + fn test_try_new_with_schema_metadata_only_replaces_metadata() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "input", + DataType::Int32, + false, + )])); + let input: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(input_schema)); + let field_metadata = + HashMap::from([("field-key".to_string(), "field-value".to_string())]); + let schema_metadata = + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]); + let metadata_schema = Schema::new_with_metadata( + vec![ + Field::new("ignored", DataType::Utf8, true) + .with_metadata(field_metadata.clone()), + ], + schema_metadata.clone(), + ); + + let projection = ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("input", 0)), + alias: "output".to_string(), + }], + input, + &metadata_schema, + )?; + + let expected_schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("output", DataType::Int32, false) + .with_metadata(field_metadata), + ], + schema_metadata, + )); + assert_eq!(projection.schema(), expected_schema); + Ok(()) + } + + fn identity_projection_with_metadata( + input: Arc<dyn ExecutionPlan>, + field_metadata: HashMap<String, String>, + schema_metadata: HashMap<String, String>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let metadata_schema = Schema::new_with_metadata( + vec![Field::new("i", DataType::Int32, true).with_metadata(field_metadata)], + schema_metadata, + ); + Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }], + input, + &metadata_schema, + )?)) + } + + #[test] + fn test_field_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert!(optimized.downcast_ref::<ProjectionExec>().is_some()); + assert_eq!(optimized.schema(), expected_schema); + Ok(()) + } + + #[test] + fn test_schema_metadata_projection_is_not_removable() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::new(), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert!(optimized.downcast_ref::<ProjectionExec>().is_some()); + assert_eq!(optimized.schema(), expected_schema); + Ok(()) + } + + #[test] + fn test_replace_children_recomputes_metadata_override() -> Result<()> { + let field_metadata = + HashMap::from([("event_field".to_string(), "true".to_string())]); + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + field_metadata.clone(), + HashMap::new(), + )?; + assert!( + projection + .downcast_ref::<ProjectionExec>() + .expect("test plan should be a ProjectionExec") + .overrides_metadata() + ); + + let replacement_schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true).with_metadata(field_metadata), + ])); + let replacement: Arc<dyn ExecutionPlan> = + Arc::new(EmptyExec::new(replacement_schema)); + let replaced = projection.replace_children( + vec![replacement], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + assert!( + !replaced + .downcast_ref::<ProjectionExec>() + .expect("replaced plan should be a ProjectionExec") + .overrides_metadata() + ); + Ok(()) + } + + #[test] + fn test_make_with_child_preserves_output_metadata() -> Result<()> { + let projection = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + let projection = projection + .downcast_ref::<ProjectionExec>() + .expect("test plan should be a ProjectionExec"); + + let rebuilt = make_with_child(projection, &test::scan_partitioned(1))?; + + assert_eq!(rebuilt.schema(), projection.schema()); + Ok(()) + } + + #[tokio::test] + async fn test_metadata_observing_parent_blocks_projection_collapse() -> Result<()> { + let inner = identity_projection_with_metadata( + test::scan_partitioned(1), + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let arrow_metadata = ScalarFunctionExpr::new( + "arrow_metadata", + Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())), + vec![ + Arc::new(Column::new("i", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "event_field".to_string(), + )))), + ], + Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + ); + let outer: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new( + [ProjectionExpr { + expr: Arc::new(arrow_metadata), + alias: "metadata".to_string(), + }], + inner, + )?); + + let outer_projection = outer + .downcast_ref::<ProjectionExec>() + .expect("test plan should be a ProjectionExec"); + assert!(try_collapse_projection_chain(outer_projection)?.is_none()); + + let optimized = remove_unnecessary_projections(outer)?.data; + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::<StringArray>() + .expect("metadata expression should return Utf8"); + assert_eq!(values.value(0), "true"); + Ok(()) + } + + #[tokio::test] + async fn test_metadata_observing_filter_blocks_projection_pushdown() -> Result<()> { + let widened: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new( + [ + ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "i".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("i", 0)), + alias: "j".to_string(), + }, + ], + test::scan_partitioned(1), + )?); + let arrow_metadata = Arc::new(ScalarFunctionExpr::new( + "arrow_metadata", + Arc::new(ScalarUDF::new_from_impl(ArrowMetadataFunc::new())), + vec![ + Arc::new(Column::new("i", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some( + "event_field".to_string(), + )))), + ], + Arc::new(Field::new("arrow_metadata", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )); + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(is_null(arrow_metadata)?, widened)?); + let projection = identity_projection_with_metadata( + filter, + HashMap::from([("event_field".to_string(), "true".to_string())]), + HashMap::new(), + )?; + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + assert_eq!(optimized.schema(), expected_schema); + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::<usize>(), + 100 + ); + Ok(()) + } + + // A schema-only metadata override must block projection embedding. The filter + // rebuilds the schema from expressions and would otherwise drop this metadata. + #[tokio::test] + async fn test_schema_level_metadata_blocks_projection_embedding() -> Result<()> { + let scan = test::scan_partitioned(1); + let predicate = binary( + col("i", &scan.schema())?, + Operator::Gt, + lit(ScalarValue::Int32(Some(-1))), + &scan.schema(), + )?; + let filter: Arc<dyn ExecutionPlan> = + Arc::new(FilterExec::try_new(predicate, scan)?); + let projection = identity_projection_with_metadata( + filter, + HashMap::new(), + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]), + )?; + // Field metadata matches, so this checks the schema-level comparison. + let projection_exec = projection + .downcast_ref::<ProjectionExec>() + .expect("test plan should be a ProjectionExec"); + assert!(projection_exec.overrides_metadata()); + let expected_schema = projection.schema(); + + let optimized = remove_unnecessary_projections(projection)?.data; + + assert_eq!(optimized.schema(), expected_schema); + assert_eq!( + optimized.schema().metadata(), + &HashMap::from([("schema-key".to_string(), "schema-value".to_string())]) + ); + + let batches = + collect(optimized.execute(0, Arc::new(TaskContext::default()))?).await?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::<usize>(), + 100 + ); + Ok(()) + } + + #[test] + fn test_collect_column_indices() -> Result<()> { + let expr = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 7)), + Operator::Minus, + Arc::new(BinaryExpr::new( + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + Operator::Plus, + Arc::new(Column::new("a", 1)), + )), + )); + let column_indices = collect_column_indices(&[ProjectionExpr { + expr, + alias: "b-(1+a)".to_string(), + }]); + // Tree traversal order: b@7 is visited before a@1 + assert_eq!(column_indices, vec![7, 1]); + Ok(()) + } + + #[test] + fn test_try_pushdown_through_join_validates_column_indices() -> Result<()> { + let child_schema = + Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let left: Arc<dyn ExecutionPlan> = + Arc::new(EmptyExec::new(Arc::clone(&child_schema))); + let right: Arc<dyn ExecutionPlan> = Arc::new(EmptyExec::new(child_schema)); + let join_schema = Arc::new(Schema::new(vec![ + Field::new("left_i", DataType::Int32, false), + Field::new("right_i", DataType::Int32, false), + ])); + let join: Arc<dyn ExecutionPlan> = + Arc::new(EmptyExec::new(Arc::clone(&join_schema))); + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("left_i", 0)), + alias: "left_i".to_string(), + }], + join, + )?; + + let Err(error) = try_pushdown_through_join_with_column_indices( + &projection, + &left, + &right, + &[], + &join_schema, + None, + &[], + ) else { + panic!("expected a mismatched mapping length to return an error"); + }; + assert!( + error.to_string().contains( + "Column index mapping has 0 entries but join schema has 2 fields" + ) + ); + + let invalid_child_index = [ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let Err(error) = try_pushdown_through_join_with_column_indices( + &projection, + &left, + &right, + &[], + &join_schema, + None, + &invalid_child_index, + ) else { + panic!("expected an invalid child index to return an error"); + }; + assert!(error.to_string().contains( + "Join output column 0 maps to left child column 1, but the child has 1 fields" + )); + + let wider_join_schema = Arc::new(Schema::new(vec![ + Field::new("left_i", DataType::Int32, false), + Field::new("right_i", DataType::Int32, false), + Field::new("extra", DataType::Int32, false), + ])); + let wider_join: Arc<dyn ExecutionPlan> = + Arc::new(EmptyExec::new(wider_join_schema)); + let out_of_mapping_projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("extra", 2)), + alias: "extra".to_string(), + }], + wider_join, + )?; + let valid_child_indices = [ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let Err(error) = try_pushdown_through_join_with_column_indices( + &out_of_mapping_projection, + &left, + &right, + &[], + &join_schema, + None, + &valid_child_indices, + ) else { + panic!("expected an out-of-mapping projection to return an error"); + }; + assert!( + error.to_string().contains( + "Projection column 2 is outside the 2-entry column index mapping" + ) + ); + + Ok(()) + } + + #[test] + fn test_join_table_borders() -> Result<()> { + let projections = vec![ + (Column::new("b", 1), "b".to_owned()), + (Column::new("c", 2), "c".to_owned()), + (Column::new("e", 4), "e".to_owned()), + (Column::new("d", 3), "d".to_owned()), + (Column::new("c", 2), "c".to_owned()), + (Column::new("f", 5), "f".to_owned()), + (Column::new("h", 7), "h".to_owned()), + (Column::new("g", 6), "g".to_owned()), + ]; + let left_table_column_count = 5; + assert_eq!( + join_table_borders(left_table_column_count, &projections), + (4, 5) + ); + + let left_table_column_count = 8; + assert_eq!( + join_table_borders(left_table_column_count, &projections), + (7, 8) + ); + + let left_table_column_count = 1; + assert_eq!( + join_table_borders(left_table_column_count, &projections), + (-1, 0) + ); + + let projections = vec![ + (Column::new("a", 0), "a".to_owned()), + (Column::new("b", 1), "b".to_owned()), + (Column::new("d", 3), "d".to_owned()), + (Column::new("g", 6), "g".to_owned()), + (Column::new("e", 4), "e".to_owned()), + (Column::new("f", 5), "f".to_owned()), + (Column::new("e", 4), "e".to_owned()), + (Column::new("h", 7), "h".to_owned()), + ]; + let left_table_column_count = 5; + assert_eq!( + join_table_borders(left_table_column_count, &projections), + (2, 7) + ); + + let left_table_column_count = 7; + assert_eq!( + join_table_borders(left_table_column_count, &projections), + (6, 7) + ); + + Ok(()) + } + + #[tokio::test] + async fn project_no_column() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + let exec = test::scan_partitioned(1); + let expected = collect(exec.execute(0, Arc::clone(&task_ctx))?).await?; + + let projection = ProjectionExec::try_new(vec![] as Vec<ProjectionExpr>, exec)?; + let stream = projection.execute(0, Arc::clone(&task_ctx))?; + let output = collect(stream).await?; + assert_eq!(output.len(), expected.len()); + + Ok(()) + } + + #[tokio::test] + async fn project_old_syntax() { + let exec = test::scan_partitioned(1); + let schema = exec.schema(); + let expr = col("i", &schema).unwrap(); + ProjectionExec::try_new( + vec![ + // use From impl of ProjectionExpr to create ProjectionExpr + // to test old syntax + (expr, "c".to_string()), + ], + exec, + ) + // expect this to succeed + .unwrap(); + } + + #[test] + fn test_projection_statistics_uses_input_schema() { + let input_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + Field::new("e", DataType::Int32, false), + Field::new("f", DataType::Int32, false), + ]); + + let input_statistics = Statistics { + num_rows: Precision::Exact(10), + column_statistics: vec![ + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(1))), + max_value: Precision::Exact(ScalarValue::Int32(Some(100))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(5))), + max_value: Precision::Exact(ScalarValue::Int32(Some(50))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(10))), + max_value: Precision::Exact(ScalarValue::Int32(Some(40))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(20))), + max_value: Precision::Exact(ScalarValue::Int32(Some(30))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(21))), + max_value: Precision::Exact(ScalarValue::Int32(Some(29))), + ..Default::default() + }, + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(24))), + max_value: Precision::Exact(ScalarValue::Int32(Some(26))), + ..Default::default() + }, + ], + ..Default::default() + }; + + let input = Arc::new(StatisticsExec::new(input_statistics, input_schema)); + + // Create projection expressions that reference columns from the input schema and the length + // of output schema columns < input schema columns and hence if we use the last few columns + // from the input schema in the expressions here, bounds_check would fail on them if output + // schema is supplied to the partitions_statistics method. + let exprs: Vec<ProjectionExpr> = vec![ + ProjectionExpr { + expr: Arc::new(Column::new("c", 2)) as Arc<dyn PhysicalExpr>, + alias: "c_renamed".to_string(), + }, + ProjectionExpr { + expr: Arc::new(BinaryExpr::new( + Arc::new(Column::new("e", 4)), + Operator::Plus, + Arc::new(Column::new("f", 5)), + )) as Arc<dyn PhysicalExpr>, + alias: "e_plus_f".to_string(), + }, + ]; + + let projection = ProjectionExec::try_new(exprs, input).unwrap(); + + let stats = StatisticsContext::new() + .compute(&projection, &StatisticsArgs::new()) + .unwrap(); + + assert_eq!(stats.num_rows, Precision::Exact(10)); + assert_eq!( + stats.column_statistics.len(), + 2, + "Expected 2 columns in projection statistics" + ); + assert!(stats.total_byte_size.is_exact().unwrap_or(false)); + } + + #[test] + fn test_filter_pushdown_with_alias() -> Result<()> { + let input_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&input_schema), + input_schema.clone(), + )); + + // project "a" as "b" + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("a", 0)), + alias: "b".to_string(), + }], + input, + )?; + + // filter "b > 5" + let filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![filter], + &ConfigOptions::default(), + )?; + + // Should be converted to "a > 5" + // "a" is index 0 in input + let expected_filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + assert_eq!(description.self_filters(), vec![vec![]]); + let pushed_filters = &description.parent_filters()[0]; + assert_eq!( + format!("{}", pushed_filters[0].predicate), + format!("{}", expected_filter) + ); + // Verify the predicate was actually pushed down + assert!(matches!(pushed_filters[0].discriminant, PushedDown::Yes)); + + Ok(()) + } + + #[test] + fn test_filter_pushdown_with_multiple_aliases() -> Result<()> { + let input_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + column_statistics: vec![Default::default(); input_schema.fields().len()], + ..Default::default() + }, + input_schema.clone(), + )); + + // project "a" as "x", "b" as "y" + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: Arc::new(Column::new("a", 0)), + alias: "x".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("b", 1)), + alias: "y".to_string(), + }, + ], + input, + )?; + + // filter "x > 5" + let filter1 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + // filter "y < 10" + let filter2 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("y", 1)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![filter1, filter2], + &ConfigOptions::default(), + )?; + + // Should be converted to "a > 5" and "b < 10" + let expected_filter1 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + let expected_filter2 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + + let pushed_filters = &description.parent_filters()[0]; + assert_eq!(pushed_filters.len(), 2); + // Note: The order of filters is preserved + assert_eq!( + format!("{}", pushed_filters[0].predicate), + format!("{}", expected_filter1) + ); + assert_eq!( + format!("{}", pushed_filters[1].predicate), + format!("{}", expected_filter2) + ); + // Verify the predicates were actually pushed down + assert!(matches!(pushed_filters[0].discriminant, PushedDown::Yes)); + assert!(matches!(pushed_filters[1].discriminant, PushedDown::Yes)); + + Ok(()) + } + + #[test] + fn test_filter_pushdown_with_swapped_aliases() -> Result<()> { + let input_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + column_statistics: vec![Default::default(); input_schema.fields().len()], + ..Default::default() + }, + input_schema.clone(), + )); + + // project "a" as "b", "b" as "a" + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: Arc::new(Column::new("a", 0)), + alias: "b".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("b", 1)), + alias: "a".to_string(), + }, + ], + input, + )?; + + // filter "b > 5" (output column 0, which is "a" in input) + let filter1 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + // filter "a < 10" (output column 1, which is "b" in input) + let filter2 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 1)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![filter1, filter2], + &ConfigOptions::default(), + )?; + + let pushed_filters = &description.parent_filters()[0]; + assert_eq!(pushed_filters.len(), 2); + + // "b" (output index 0) -> "a" (input index 0) + let expected_filter1 = "a@0 > 5"; + // "a" (output index 1) -> "b" (input index 1) + let expected_filter2 = "b@1 < 10"; + + assert_eq!(format!("{}", pushed_filters[0].predicate), expected_filter1); + assert_eq!(format!("{}", pushed_filters[1].predicate), expected_filter2); + // Verify the predicates were actually pushed down + assert!(matches!(pushed_filters[0].discriminant, PushedDown::Yes)); + assert!(matches!(pushed_filters[1].discriminant, PushedDown::Yes)); + + Ok(()) + } + + #[test] + fn test_filter_pushdown_with_mixed_columns() -> Result<()> { + let input_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let input = Arc::new(StatisticsExec::new( + Statistics { + column_statistics: vec![Default::default(); input_schema.fields().len()], + ..Default::default() + }, + input_schema.clone(), + )); + + // project "a" as "x", "b" as "b" (pass through) + let projection = ProjectionExec::try_new( + vec![ + ProjectionExpr { + expr: Arc::new(Column::new("a", 0)), + alias: "x".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("b", 1)), + alias: "b".to_string(), + }, + ], + input, + )?; + + // filter "x > 5" + let filter1 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("x", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + // filter "b < 10" (using output index 1 which corresponds to 'b') + let filter2 = Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![filter1, filter2], + &ConfigOptions::default(), + )?; + + let pushed_filters = &description.parent_filters()[0]; + assert_eq!(pushed_filters.len(), 2); + // "x" -> "a" (index 0) + let expected_filter1 = "a@0 > 5"; + // "b" -> "b" (index 1) + let expected_filter2 = "b@1 < 10"; + + assert_eq!(format!("{}", pushed_filters[0].predicate), expected_filter1); + assert_eq!(format!("{}", pushed_filters[1].predicate), expected_filter2); + // Verify the predicates were actually pushed down + assert!(matches!(pushed_filters[0].discriminant, PushedDown::Yes)); + assert!(matches!(pushed_filters[1].discriminant, PushedDown::Yes)); + + Ok(()) + } + + #[test] + fn test_filter_pushdown_with_complex_expression() -> Result<()> { + let input_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + column_statistics: vec![Default::default(); input_schema.fields().len()], + ..Default::default() + }, + input_schema.clone(), + )); + + // project "a + 1" as "z" + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )), + alias: "z".to_string(), + }], + input, + )?; + + // filter "z > 10" + let filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("z", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(10)))), + )) as Arc<dyn PhysicalExpr>; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![filter], + &ConfigOptions::default(), + )?; + + // expand to `a + 1 > 10` + let pushed_filters = &description.parent_filters()[0]; + assert!(matches!(pushed_filters[0].discriminant, PushedDown::Yes)); + assert_eq!(format!("{}", pushed_filters[0].predicate), "a@0 + 1 > 10"); + + Ok(()) + } + + #[test] + fn test_filter_pushdown_with_unknown_column() -> Result<()> { + let input_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + column_statistics: vec![Default::default(); input_schema.fields().len()], + ..Default::default() + }, + input_schema.clone(), + )); + + // project "a" as "a" + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("a", 0)), + alias: "a".to_string(), + }], + input, + )?; + + // filter "unknown_col > 5" - using a column name that doesn't exist in projection output + // Column constructor: name, index. Index 1 doesn't exist. + let filter = Arc::new(BinaryExpr::new( + Arc::new(Column::new("unknown_col", 1)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )) as Arc<dyn PhysicalExpr>; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![filter], + &ConfigOptions::default(), + )?; + + let pushed_filters = &description.parent_filters()[0]; + assert!(matches!(pushed_filters[0].discriminant, PushedDown::No)); + // The column shouldn't be found in the alias map, so it remains unchanged with its index + assert_eq!( + format!("{}", pushed_filters[0].predicate), + "unknown_col@1 > 5" + ); + + Ok(()) + } + + /// Basic test for `DynamicFilterPhysicalExpr` can correctly update its child expression + /// i.e. starting with lit(true) and after update it becomes `a > 5` + /// with projection [b - 1 as a], the pushed down filter should be `b - 1 > 5` + #[test] + fn test_basic_dyn_filter_projection_pushdown_update_child() -> Result<()> { + let input_schema = + Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); + + let input = Arc::new(StatisticsExec::new( + Statistics { + column_statistics: vec![Default::default(); input_schema.fields().len()], + ..Default::default() + }, + input_schema.as_ref().clone(), + )); + + // project "b" - 1 as "a" + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: binary( + Arc::new(Column::new("b", 0)), + Operator::Minus, + lit(1), + &input_schema, + ) + .unwrap(), + alias: "a".to_string(), + }], + input, + )?; + + // simulate projection's parent create a dynamic filter on "a" + let projected_schema = projection.schema(); + let col_a = col("a", &projected_schema)?; + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + lit(true), + )); + // Initial state should be lit(true) + let current = dynamic_filter.current()?; + assert_eq!(format!("{current}"), "true"); + + let dyn_phy_expr: Arc<dyn PhysicalExpr> = Arc::clone(&dynamic_filter) as _; + + let description = projection.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![dyn_phy_expr], + &ConfigOptions::default(), + )?; + + let pushed_filters = &description.parent_filters()[0][0]; + + // Check currently pushed_filters is lit(true) + assert_eq!( + format!("{}", pushed_filters.predicate), + "DynamicFilter [ empty ]" + ); + + // Update to a > 5 (after projection, b is now called a) + let new_expr = + Arc::new(BinaryExpr::new(Arc::clone(&col_a), Operator::Gt, lit(5i32))); + dynamic_filter.update(new_expr)?; + + // Now it should be a > 5 + let current = dynamic_filter.current()?; + assert_eq!(format!("{current}"), "a@0 > 5"); + + // Check currently pushed_filters is b - 1 > 5 (because b - 1 is projected as a) + assert_eq!( + format!("{}", pushed_filters.predicate), + "DynamicFilter [ b@0 - 1 > 5 ]" + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/proto.rs b/native/vendor/datafusion-physical-plan/src/proto.rs new file mode 100644 index 0000000..7640d76 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/proto.rs
@@ -0,0 +1,386 @@ +// 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. + +//! Serialization hooks for [`ExecutionPlan`], mirroring the +//! `try_to_proto`/`try_from_proto` pattern used for `PhysicalExpr`. +//! +//! # Why the indirection +//! +//! An `ExecutionPlan` must be able to (de)serialize its child plans and its +//! child physical expressions recursively. The concrete recursion lives in +//! `datafusion-proto` (it owns the extension codec, the session context and the +//! central converter), but `datafusion-proto` sits *above* `datafusion-physical-plan` +//! in the crate graph. To let a plan drive that recursion without a dependency +//! cycle, this module defines: +//! +//! * [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] — the stable, +//! concrete context types a plan author interacts with. New capabilities can +//! be added here without changing every plan's hook signature. +//! * [`ExecutionPlanEncode`] / [`ExecutionPlanDecode`] — internal dispatch +//! traits, *defined* here but *implemented* in `datafusion-proto`, that the +//! context types delegate to. This is the dependency inversion that keeps the +//! proto types flowing in one direction only. They are `#[doc(hidden)]`: not +//! public API, `pub` only because their implementors live in another crate. +//! +//! `datafusion-physical-plan` depends on the pure prost types in +//! `datafusion-proto-models` (feature `proto`), never on `datafusion-proto`. +//! +//! # Function-carrying plans +//! +//! Plans that reference UD(A/W)Fs (`AggregateExec`, the window execs, …) also +//! ride the hook: the context exposes typed, *bytes-only* function serde — +//! [`encode_udaf`](ExecutionPlanEncodeCtx::encode_udaf) / +//! [`decode_udaf`](ExecutionPlanDecodeCtx::decode_udaf) and the udf/udwf +//! siblings. These take/return `datafusion-expr` types plus `Vec<u8>` and never +//! name a proto type, so the `PhysicalExtensionCodec` (which only +//! `datafusion-proto` can name) stays fully encapsulated behind the adapter that +//! backs these traits. The lookup-order policy (payload → codec; else registry → +//! codec fallback) lives once, in that adapter, rather than in every plan. +//! +//! This is possible because `datafusion-physical-plan` sits *above* +//! `datafusion-expr` in the crate graph; the expression-side ctx (in +//! `physical-expr-common`, *below* `datafusion-expr`) cannot do this, which is +//! why `ScalarFunctionExpr` remains special-cased there. +//! +//! [`ExecutionPlan`]: crate::ExecutionPlan + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::physical_planning_context::ScalarSubqueryResults; +use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, +}; +use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, +}; +use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; + +use crate::ExecutionPlan; + +/// Internal dispatch trait backing [`ExecutionPlanEncodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanEncodeCtx`] instead. +/// +/// **Not public API.** `pub` only because the implementors live in another +/// crate; `#[doc(hidden)]` records that, so encoding primitives can be added +/// here as the serialization hooks grow without breaking downstream code. +#[doc(hidden)] +pub trait ExecutionPlanEncode { + /// Serialize a child execution plan (recursing through the central + /// serializer, so the child's own `try_to_proto` hook is honored). + fn encode_plan(&self, plan: &Arc<dyn ExecutionPlan>) -> Result<PhysicalPlanNode>; + + /// Serialize a physical expression owned by the plan. + fn encode_expr(&self, expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode>; + + /// Serialize a scalar UDF to an opaque payload. `None` means "decodable by + /// name alone" (built-ins). Bytes-only: no proto types cross this boundary. + fn encode_udf(&self, udf: &ScalarUDF) -> Result<Option<Vec<u8>>>; + + /// Serialize an aggregate UDF to an opaque payload. `None` means "decodable + /// by name alone". + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result<Option<Vec<u8>>>; + + /// Serialize a window UDF to an opaque payload. `None` means "decodable by + /// name alone". + fn encode_udwf(&self, udwf: &WindowUDF) -> Result<Option<Vec<u8>>>; +} + +/// Internal dispatch trait backing [`ExecutionPlanDecodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanDecodeCtx`] instead. +/// +/// **Not public API.** `pub` only because the implementors live in another +/// crate; `#[doc(hidden)]` records that, so decoding primitives can be added +/// here as the serialization hooks grow without breaking downstream code. +#[doc(hidden)] +pub trait ExecutionPlanDecode { + /// Deserialize a child execution plan (recursing through the central + /// deserializer, so the child's own `try_from_proto` is honored). + fn decode_plan(&self, node: &PhysicalPlanNode) -> Result<Arc<dyn ExecutionPlan>>; + + /// Deserialize a child plan with `results` active for scalar subquery + /// expressions in that plan's subtree. + fn decode_plan_with_scalar_subquery_results( + &self, + node: &PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result<Arc<dyn ExecutionPlan>>; + + /// Deserialize a physical expression against `input_schema`. + fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>>; + + /// The session task context, used by plans that need the function registry + /// or session configuration. Never exposes the proto extension codec. + fn task_ctx(&self) -> &TaskContext; + + /// Reconstruct a scalar UDF from its name and optional payload. Encapsulates + /// the lookup-order policy (payload → codec; else registry → codec fallback) + /// so no plan re-derives it. Bytes-only: no proto types cross this boundary. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result<Arc<ScalarUDF>>; + + /// Reconstruct an aggregate UDF from its name and optional payload. + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result<Arc<AggregateUDF>>; + + /// Reconstruct a window UDF from its name and optional payload. + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result<Arc<WindowUDF>>; +} + +/// Context handed to [`ExecutionPlan::try_to_proto`]. +/// +/// +/// Provides the primitives a plan needs to serialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanEncodeCtx<'a> { + encoder: &'a dyn ExecutionPlanEncode, +} + +impl<'a> ExecutionPlanEncodeCtx<'a> { + /// Create a new encode context wrapping an [`ExecutionPlanEncode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(encoder: &'a dyn ExecutionPlanEncode) -> Self { + Self { encoder } + } + + /// Serialize a single child plan. + pub fn encode_child( + &self, + plan: &Arc<dyn ExecutionPlan>, + ) -> Result<PhysicalPlanNode> { + self.encoder.encode_plan(plan) + } + + /// Serialize an iterator of child plans. + pub fn encode_children<'b, I>(&self, plans: I) -> Result<Vec<PhysicalPlanNode>> + where + I: IntoIterator<Item = &'b Arc<dyn ExecutionPlan>>, + { + plans.into_iter().map(|p| self.encode_child(p)).collect() + } + + /// Serialize a single physical expression. + pub fn encode_expr(&self, expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> { + self.encoder.encode_expr(expr) + } + + /// Serialize an iterator of physical expressions. + pub fn encode_expressions<'b, I>(&self, exprs: I) -> Result<Vec<PhysicalExprNode>> + where + I: IntoIterator<Item = &'b Arc<dyn PhysicalExpr>>, + { + exprs.into_iter().map(|e| self.encode_expr(e)).collect() + } + + /// Serialize a scalar UDF to an opaque payload (`None` = built-in, decodable + /// by name). No proto types cross this boundary. + pub fn encode_udf(&self, udf: &ScalarUDF) -> Result<Option<Vec<u8>>> { + self.encoder.encode_udf(udf) + } + + /// Serialize an aggregate UDF to an opaque payload (`None` = decodable by + /// name). + pub fn encode_udaf(&self, udaf: &AggregateUDF) -> Result<Option<Vec<u8>>> { + self.encoder.encode_udaf(udaf) + } + + /// Serialize a window UDF to an opaque payload (`None` = decodable by name). + pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result<Option<Vec<u8>>> { + self.encoder.encode_udwf(udwf) + } + + /// An expression-level encode context backed by this plan context. + /// + /// Lets a plan hand `ctx` to expression-level conversions that own their own + /// wire logic — e.g. + /// [`Partitioning::try_to_proto`](datafusion_physical_expr::Partitioning::try_to_proto) + /// and + /// [`PhysicalSortExpr::try_to_proto`](datafusion_physical_expr::PhysicalSortExpr::try_to_proto). + pub fn expr_ctx(&self) -> PhysicalExprEncodeCtx<'_> { + PhysicalExprEncodeCtx::new(self) + } +} + +/// Lets [`ExecutionPlanEncodeCtx`] back a [`PhysicalExprEncodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprEncode for ExecutionPlanEncodeCtx<'_> { + fn encode(&self, expr: &Arc<dyn PhysicalExpr>) -> Result<PhysicalExprNode> { + self.encode_expr(expr) + } +} + +/// Context handed to a plan's `try_from_proto` associated function. +/// +/// Provides the primitives a plan needs to deserialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanDecodeCtx<'a> { + decoder: &'a dyn ExecutionPlanDecode, +} + +impl<'a> ExecutionPlanDecodeCtx<'a> { + /// Create a new decode context wrapping an [`ExecutionPlanDecode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(decoder: &'a dyn ExecutionPlanDecode) -> Self { + Self { decoder } + } + + /// Deserialize a single child plan. + pub fn decode_child( + &self, + node: &PhysicalPlanNode, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.decoder.decode_plan(node) + } + + /// Deserialize a child plan with `results` active for scalar subquery + /// expressions in that plan's subtree. + pub fn decode_child_with_scalar_subquery_results( + &self, + node: &PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.decoder + .decode_plan_with_scalar_subquery_results(node, results) + } + + /// Deserialize a required child plan, producing a uniform "missing required + /// field" error when the optional wire field is absent. + pub fn decode_required_child( + &self, + node: Option<&PhysicalPlanNode>, + plan_name: &str, + field: &str, + ) -> Result<Arc<dyn ExecutionPlan>> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_child(node) + } + + /// Deserialize a physical expression against `input_schema`. + pub fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.decoder.decode_expr(node, input_schema) + } + + /// Deserialize a required physical expression against `input_schema`. + pub fn decode_required_expr( + &self, + node: Option<&PhysicalExprNode>, + input_schema: &Schema, + plan_name: &str, + field: &str, + ) -> Result<Arc<dyn PhysicalExpr>> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_expr(node, input_schema) + } + + /// The session task context (function registry + session config). Never + /// exposes the proto extension codec. + pub fn task_ctx(&self) -> &TaskContext { + self.decoder.task_ctx() + } + + /// Reconstruct a scalar UDF from its name and optional payload. The + /// lookup-order policy is owned by `datafusion-proto`; no proto types cross + /// this boundary. + pub fn decode_udf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result<Arc<ScalarUDF>> { + self.decoder.decode_udf(name, payload) + } + + /// Reconstruct an aggregate UDF from its name and optional payload. + pub fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result<Arc<AggregateUDF>> { + self.decoder.decode_udaf(name, payload) + } + + /// Reconstruct a window UDF from its name and optional payload. + pub fn decode_udwf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result<Arc<WindowUDF>> { + self.decoder.decode_udwf(name, payload) + } + + /// An expression-level decode context backed by this plan context, bound to + /// `input_schema`. + /// + /// The decode counterpart of + /// [`ExecutionPlanEncodeCtx::expr_ctx`], for calling conversions such as + /// [`Partitioning::try_from_proto`](datafusion_physical_expr::Partitioning::try_from_proto). + pub fn expr_ctx<'s>(&'s self, input_schema: &'s Schema) -> PhysicalExprDecodeCtx<'s> { + PhysicalExprDecodeCtx::new(input_schema, self) + } +} + +/// Lets [`ExecutionPlanDecodeCtx`] back a [`PhysicalExprDecodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> { + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.decode_expr(node, schema) + } +} + +/// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` +/// variant, returning a reference to the inner payload, else an `internal_err!`. +/// Mirrors `expect_expr_variant!` on the expression side. Field access on the +/// result auto-derefs through the `Box` that boxed variants use. +#[macro_export] +macro_rules! expect_plan_variant { + ($node:expr, $variant:path, $plan_name:literal $(,)?) => {{ + match &$node.physical_plan_type { + Some($variant(inner)) => inner, + _ => { + return ::datafusion_common::internal_err!(concat!( + "PhysicalPlanNode is not a ", + $plan_name + )); + } + } + }}; +}
diff --git a/native/vendor/datafusion-physical-plan/src/recursive_query.rs b/native/vendor/datafusion-physical-plan/src/recursive_query.rs new file mode 100644 index 0000000..0a56488 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/recursive_query.rs
@@ -0,0 +1,581 @@ +// 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. + +//! Defines the recursive query plan + +use std::any::Any; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::work_table::{ReservedBatches, WorkTable}; +use crate::aggregates::group_values::{GroupValues, new_group_values}; +use crate::aggregates::order::GroupOrdering; +use crate::common::project_plan_to_schema; +use crate::execution_plan::{Boundedness, EmissionType, reset_plan_states}; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, +}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use arrow::array::{BooleanArray, BooleanBuilder}; +use arrow::compute::filter_record_batch; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::{ + Result, exec_datafusion_err, internal_datafusion_err, not_impl_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; + +use futures::{Stream, StreamExt, ready}; + +/// Recursive query execution plan. +/// +/// This plan has two components: a base part (the static term) and +/// a dynamic part (the recursive term). The execution will start from +/// the base, and as long as the previous iteration produced at least +/// a single new row (taking care of the distinction) the recursive +/// part will be continuously executed. +/// +/// Before each execution of the dynamic part, the rows from the previous +/// iteration will be available in a "working table" (not a real table, +/// can be only accessed using a continuance operation). +/// +/// Note that there won't be any limit or checks applied to detect +/// an infinite recursion, so it is up to the planner to ensure that +/// it won't happen. +#[derive(Debug, Clone)] +pub struct RecursiveQueryExec { + /// Name of the query handler + name: String, + /// The working table of cte + work_table: Arc<WorkTable>, + /// The base part (static term) + static_term: Arc<dyn ExecutionPlan>, + /// The dynamic part (recursive term) + recursive_term: Arc<dyn ExecutionPlan>, + /// Distinction + is_distinct: bool, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl RecursiveQueryExec { + /// Create a new RecursiveQueryExec + pub fn try_new( + name: String, + output_schema: SchemaRef, + static_term: Arc<dyn ExecutionPlan>, + recursive_term: Arc<dyn ExecutionPlan>, + is_distinct: bool, + ) -> Result<Self> { + // Each recursive query needs its own work table + let work_table = Arc::new(WorkTable::new(name.clone())); + // Use the same work table for both the WorkTableExec and the recursive term + let static_term = project_plan_to_schema(static_term, &output_schema)?; + let recursive_term = assign_work_table(recursive_term, &work_table)?; + let recursive_term = project_plan_to_schema(recursive_term, &output_schema)?; + let cache = Self::compute_properties(output_schema); + Ok(RecursiveQueryExec { + name, + static_term, + recursive_term, + is_distinct, + work_table, + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + }) + } + + /// Ref to name + pub fn name(&self) -> &str { + &self.name + } + + /// Ref to static term + pub fn static_term(&self) -> &Arc<dyn ExecutionPlan> { + &self.static_term + } + + /// Ref to recursive term + pub fn recursive_term(&self) -> &Arc<dyn ExecutionPlan> { + &self.recursive_term + } + + /// is distinct + pub fn is_distinct(&self) -> bool { + self.is_distinct + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + let eq_properties = EquivalenceProperties::new(schema); + + PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl ExecutionPlan for RecursiveQueryExec { + fn name(&self) -> &'static str { + "RecursiveQueryExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.static_term, &self.recursive_term] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + // TODO: control these hints and see whether we can + // infer some from the child plans (static/recursive terms). + fn maintains_input_order(&self) -> Vec<bool> { + vec![false, false] + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false, false] + } + + fn required_input_distribution(&self) -> Vec<crate::Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + crate::Distribution::SinglePartition, + crate::Distribution::SinglePartition, + ]) + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + RecursiveQueryExec::try_new( + self.name.clone(), + self.schema(), + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.is_distinct, + ) + .map(|e| Arc::new(e) as _) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + // TODO: we might be able to handle multiple partitions in the future. + if partition != 0 { + return Err(internal_datafusion_err!( + "RecursiveQueryExec got an invalid partition {partition} (expected 0)" + )); + } + + let static_stream = self.static_term.execute(partition, Arc::clone(&context))?; + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + Ok(Box::pin(RecursiveQueryStream::new( + context, + Arc::clone(&self.work_table), + Arc::clone(&self.recursive_term), + static_stream, + self.is_distinct, + baseline_metrics, + )?)) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } +} + +impl DisplayAs for RecursiveQueryExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "RecursiveQueryExec: name={}, is_distinct={}", + self.name, self.is_distinct + ) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +/// The actual logic of the recursive queries happens during the streaming +/// process. A simplified version of the algorithm is the following: +/// +/// buffer = [] +/// +/// while batch := static_stream.next(): +/// buffer.push(batch) +/// yield buffer +/// +/// while buffer.len() > 0: +/// sender, receiver = Channel() +/// register_continuation(handle_name, receiver) +/// sender.send(buffer.drain()) +/// recursive_stream = recursive_term.execute() +/// while batch := recursive_stream.next(): +/// buffer.append(batch) +/// yield buffer +struct RecursiveQueryStream { + /// The context to be used for managing handlers & executing new tasks + task_context: Arc<TaskContext>, + /// The working table state, representing the self referencing cte table + work_table: Arc<WorkTable>, + /// The dynamic part (recursive term) as is (without being executed) + recursive_term: Arc<dyn ExecutionPlan>, + /// The static part (static term) as a stream. If the processing of this + /// part is completed, then it will be None. + static_stream: Option<SendableRecordBatchStream>, + /// The dynamic part (recursive term) as a stream. If the processing of this + /// part has not started yet, or has been completed, then it will be None. + recursive_stream: Option<SendableRecordBatchStream>, + /// The schema of the output. + schema: SchemaRef, + /// In-memory buffer for storing a copy of the current results. Will be + /// cleared after each iteration. + buffer: Vec<RecordBatch>, + /// Tracks the memory used by the buffer + reservation: MemoryReservation, + /// If the distinct flag is set, then we use this hash table to remove duplicates from result and work tables + distinct_deduplicator: Option<DistinctDeduplicator>, + /// Metrics. + baseline_metrics: BaselineMetrics, +} + +impl RecursiveQueryStream { + /// Create a new recursive query stream + fn new( + task_context: Arc<TaskContext>, + work_table: Arc<WorkTable>, + recursive_term: Arc<dyn ExecutionPlan>, + static_stream: SendableRecordBatchStream, + is_distinct: bool, + baseline_metrics: BaselineMetrics, + ) -> Result<Self> { + let schema = static_stream.schema(); + let reservation = + MemoryConsumer::new("RecursiveQuery").register(task_context.memory_pool()); + let distinct_deduplicator = is_distinct + .then(|| DistinctDeduplicator::new(Arc::clone(&schema), &task_context)) + .transpose()?; + Ok(Self { + task_context, + work_table, + recursive_term, + static_stream: Some(static_stream), + recursive_stream: None, + schema, + buffer: vec![], + reservation, + distinct_deduplicator, + baseline_metrics, + }) + } + + /// Push a clone of the given batch to the in memory buffer, and then return + /// a poll with it. + fn push_batch( + mut self: std::pin::Pin<&mut Self>, + mut batch: RecordBatch, + ) -> Poll<Option<Result<RecordBatch>>> { + let baseline_metrics = self.baseline_metrics.clone(); + + if let Some(deduplicator) = &mut self.distinct_deduplicator { + let _timer_guard = baseline_metrics.elapsed_compute().timer(); + batch = deduplicator.deduplicate(&batch)?; + } + + if let Err(e) = self.reservation.try_grow(batch.get_array_memory_size()) { + return Poll::Ready(Some(Err(e))); + } + self.buffer.push(batch.clone()); + (&batch).record_output(&baseline_metrics); + Poll::Ready(Some(Ok(batch))) + } + + /// Start polling for the next iteration, will be called either after the static term + /// is completed or another term is completed. It will follow the algorithm above on + /// to check whether the recursion has ended. + fn poll_next_iteration( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + let total_length = self + .buffer + .iter() + .fold(0, |acc, batch| acc + batch.num_rows()); + + if total_length == 0 { + return Poll::Ready(None); + } + + // Update the work table with the current buffer + let reserved_batches = ReservedBatches::new( + std::mem::take(&mut self.buffer), + self.reservation.take(), + ); + self.work_table.update(reserved_batches); + + // We always execute (and re-execute iteratively) the first partition. + // Downstream plans should not expect any partitioning. + let partition = 0; + + let recursive_plan = reset_plan_states(Arc::clone(&self.recursive_term))?; + self.recursive_stream = + Some(recursive_plan.execute(partition, Arc::clone(&self.task_context))?); + self.poll_next(cx) + } +} + +fn assign_work_table( + plan: Arc<dyn ExecutionPlan>, + work_table: &Arc<WorkTable>, +) -> Result<Arc<dyn ExecutionPlan>> { + let mut work_table_refs = 0; + plan.transform_down(|plan| { + if let Some(new_plan) = + plan.with_new_state(Arc::clone(work_table) as Arc<dyn Any + Send + Sync>) + { + if work_table_refs > 0 { + not_impl_err!( + "Multiple recursive references to the same CTE are not supported" + ) + } else { + work_table_refs += 1; + Ok(Transformed::yes(new_plan)) + } + } else { + Ok(Transformed::no(plan)) + } + }) + .data() +} + +impl Stream for RecursiveQueryStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if let Some(static_stream) = &mut self.static_stream { + // While the static term's stream is available, we'll be forwarding the batches from it (also + // saving them for the initial iteration of the recursive term). + let batch_result = ready!(static_stream.poll_next_unpin(cx)); + match &batch_result { + None => { + // Once this is done, we can start running the setup for the recursive term. + self.static_stream = None; + self.poll_next_iteration(cx) + } + Some(Ok(batch)) => self.push_batch(batch.clone()), + _ => Poll::Ready(batch_result), + } + } else if let Some(recursive_stream) = &mut self.recursive_stream { + let batch_result = ready!(recursive_stream.poll_next_unpin(cx)); + match batch_result { + None => { + self.recursive_stream = None; + self.poll_next_iteration(cx) + } + Some(Ok(batch)) => self.push_batch(batch), + _ => Poll::Ready(batch_result), + } + } else { + Poll::Ready(None) + } + } +} + +impl RecordBatchStream for RecursiveQueryStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Deduplicator based on a hash table. +struct DistinctDeduplicator { + /// Grouped rows used for distinct + group_values: Box<dyn GroupValues>, + reservation: MemoryReservation, + intern_output_buffer: Vec<usize>, +} + +impl DistinctDeduplicator { + fn new(schema: SchemaRef, task_context: &TaskContext) -> Result<Self> { + let group_values = new_group_values(schema, &GroupOrdering::None)?; + let reservation = MemoryConsumer::new("RecursiveQueryHashTable") + .register(task_context.memory_pool()); + Ok(Self { + group_values, + reservation, + intern_output_buffer: Vec::new(), + }) + } + + /// Remove duplicated rows from the given batch, keeping a state between batches. + /// + /// We use a hash table to allocate new group ids for the new rows. + /// [`GroupValues`] allocate increasing group ids. + /// Hence, if groups (i.e., rows) are new, then they have ids >= length before interning, we keep them. + /// We also detect duplicates by enforcing that group ids are increasing. + fn deduplicate(&mut self, batch: &RecordBatch) -> Result<RecordBatch> { + let size_before = self.group_values.len(); + let additional = batch.num_rows(); + self.intern_output_buffer + .try_reserve(additional) + .map_err(|e| { + exec_datafusion_err!( + "failed to reserve {additional} recursive query group ids: {e}" + ) + })?; + self.group_values + .intern(batch.columns(), &mut self.intern_output_buffer)?; + let mask = new_groups_mask(&self.intern_output_buffer, size_before); + self.intern_output_buffer.clear(); + // We update the reservation to reflect the new size of the hash table. + self.reservation.try_resize(self.group_values.size())?; + Ok(filter_record_batch(batch, &mask)?) + } +} + +/// Return a mask, each element being true if, and only if, the element is greater than all previous elements and greater or equal than the provided max_already_seen_group_id +fn new_groups_mask( + values: &[usize], + mut max_already_seen_group_id: usize, +) -> BooleanArray { + let mut output = BooleanBuilder::with_capacity(values.len()); + for value in values { + if *value >= max_already_seen_group_id { + output.append_value(true); + max_already_seen_group_id = *value + 1; // We want to be increasing + } else { + output.append_value(false); + } + } + output.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExec; + + use arrow::datatypes::{DataType, Field, Schema}; + + fn empty_exec(fields: Vec<Field>) -> Arc<dyn ExecutionPlan> { + Arc::new(EmptyExec::new(Arc::new(Schema::new(fields)))) + } + + #[test] + fn recursive_query_exec_projects_recursive_term_to_reconciled_schema() -> Result<()> { + let static_term = empty_exec(vec![Field::new("value", DataType::Int32, false)]); + let recursive_term = + empty_exec(vec![Field::new("value + Int32(1)", DataType::Int32, false)]); + + let exec = RecursiveQueryExec::try_new( + "numbers".to_string(), + static_term.schema(), + Arc::clone(&static_term), + Arc::clone(&recursive_term), + false, + )?; + + assert_eq!(exec.schema(), static_term.schema()); + let projection = exec + .recursive_term() + .downcast_ref::<ProjectionExec>() + .expect("recursive term should be aligned with ProjectionExec"); + assert!(Arc::ptr_eq(projection.input(), &recursive_term)); + assert!(!projection.schema().field(0).is_nullable()); + assert_eq!(projection.expr()[0].alias, "value"); + Ok(()) + } + + #[test] + fn recursive_query_exec_reconciles_nullability() -> Result<()> { + let static_term = empty_exec(vec![Field::new("value", DataType::Int32, false)]); + let recursive_term = + empty_exec(vec![Field::new("value + Int32(1)", DataType::Int32, true)]); + let output_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + true, + )])); + + let exec = RecursiveQueryExec::try_new( + "numbers".to_string(), + Arc::clone(&output_schema), + static_term, + recursive_term, + false, + )?; + + assert!(exec.schema().field(0).is_nullable()); + assert!(exec.static_term().schema().field(0).is_nullable()); + assert!(exec.recursive_term().schema().field(0).is_nullable()); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/render_tree.rs b/native/vendor/datafusion-physical-plan/src/render_tree.rs new file mode 100644 index 0000000..40e2763 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/render_tree.rs
@@ -0,0 +1,231 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// This code is based on the DuckDB’s implementation: +// <https://github.com/duckdb/duckdb/blob/main/src/include/duckdb/common/render_tree.hpp> + +//! This module provides functionality for rendering an execution plan as a tree structure. +//! It helps in visualizing how different operations in a query are connected and organized. + +use std::collections::HashMap; +use std::fmt::Formatter; +use std::sync::Arc; +use std::{cmp, fmt}; + +use crate::{DisplayFormatType, ExecutionPlan}; + +// TODO: It's never used. +/// Represents a 2D coordinate in the rendered tree. +/// Used to track positions of nodes and their connections. +pub struct Coordinate { + /// Horizontal position in the tree + #[expect(dead_code)] + pub x: usize, + /// Vertical position in the tree + #[expect(dead_code)] + pub y: usize, +} + +impl Coordinate { + pub fn new(x: usize, y: usize) -> Self { + Coordinate { x, y } + } +} + +/// Represents a node in the render tree, containing information about an execution plan operator +/// and its relationships to other operators. +pub struct RenderTreeNode { + /// The name of physical `ExecutionPlan`. + pub name: String, + /// Execution info collected from `ExecutionPlan`. + pub extra_text: HashMap<String, String>, + /// Positions of child nodes in the rendered tree. + pub child_positions: Vec<Coordinate>, +} + +impl RenderTreeNode { + pub fn new(name: String, extra_text: HashMap<String, String>) -> Self { + RenderTreeNode { + name, + extra_text, + child_positions: vec![], + } + } + + fn add_child_position(&mut self, x: usize, y: usize) { + self.child_positions.push(Coordinate::new(x, y)); + } +} + +/// Main structure for rendering an execution plan as a tree. +/// Manages a 2D grid of nodes and their layout information. +pub struct RenderTree { + /// Storage for tree nodes in a flattened 2D grid + pub nodes: Vec<Option<Arc<RenderTreeNode>>>, + /// Total width of the rendered tree + pub width: usize, + /// Total height of the rendered tree + pub height: usize, +} + +impl RenderTree { + /// Creates a new render tree from an execution plan. + pub fn create_tree(plan: &dyn ExecutionPlan) -> Self { + let (width, height) = get_tree_width_height(plan); + + let mut result = Self::new(width, height); + + create_tree_recursive(&mut result, plan, 0, 0); + + result + } + + fn new(width: usize, height: usize) -> Self { + RenderTree { + nodes: vec![None; (width + 1) * (height + 1)], + width, + height, + } + } + + pub fn get_node(&self, x: usize, y: usize) -> Option<Arc<RenderTreeNode>> { + if x >= self.width || y >= self.height { + return None; + } + + let pos = self.get_position(x, y); + self.nodes.get(pos).and_then(|node| node.clone()) + } + + pub fn set_node(&mut self, x: usize, y: usize, node: Arc<RenderTreeNode>) { + let pos = self.get_position(x, y); + if let Some(slot) = self.nodes.get_mut(pos) { + *slot = Some(node); + } + } + + pub fn has_node(&self, x: usize, y: usize) -> bool { + if x >= self.width || y >= self.height { + return false; + } + + let pos = self.get_position(x, y); + self.nodes.get(pos).is_some_and(|node| node.is_some()) + } + + fn get_position(&self, x: usize, y: usize) -> usize { + y * self.width + x + } +} + +/// Calculates the required dimensions of the tree. +/// This ensures we allocate enough space for the entire tree structure. +/// +/// # Arguments +/// * `plan` - The execution plan to measure +/// +/// # Returns +/// * A tuple of (width, height) representing the dimensions needed for the tree +fn get_tree_width_height(plan: &dyn ExecutionPlan) -> (usize, usize) { + let children = plan.children(); + + // Leaf nodes take up 1x1 space + if children.is_empty() { + return (1, 1); + } + + let mut width = 0; + let mut height = 0; + + for child in children { + let (child_width, child_height) = get_tree_width_height(child.as_ref()); + width += child_width; + height = cmp::max(height, child_height); + } + + height += 1; + + (width, height) +} + +fn fmt_display(plan: &dyn ExecutionPlan) -> impl fmt::Display + '_ { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + } + + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + self.plan.fmt_as(DisplayFormatType::TreeRender, f)?; + Ok(()) + } + } + + Wrapper { plan } +} + +/// Recursively builds the render tree structure. +/// Traverses the execution plan and creates corresponding render nodes while +/// maintaining proper positioning and parent-child relationships. +/// +/// # Arguments +/// * `result` - The render tree being constructed +/// * `plan` - Current execution plan node being processed +/// * `x` - Horizontal position in the tree +/// * `y` - Vertical position in the tree +/// +/// # Returns +/// * The width of the subtree rooted at the current node +fn create_tree_recursive( + result: &mut RenderTree, + plan: &dyn ExecutionPlan, + x: usize, + y: usize, +) -> usize { + let display_info = fmt_display(plan).to_string(); + let mut extra_info = HashMap::new(); + + // Parse the key-value pairs from the formatted string. + // See DisplayFormatType::TreeRender for details + for line in display_info.lines() { + if let Some((key, value)) = line.split_once('=') { + extra_info.insert(key.to_string(), value.to_string()); + } else { + extra_info.insert(line.to_string(), "".to_string()); + } + } + + let mut node = RenderTreeNode::new(plan.name().to_string(), extra_info); + + let children = plan.children(); + + if children.is_empty() { + result.set_node(x, y, Arc::new(node)); + return 1; + } + + let mut width = 0; + for child in children { + let child_x = x + width; + let child_y = y + 1; + node.add_child_position(child_x, child_y); + width += create_tree_recursive(result, child.as_ref(), child_x, child_y); + } + + result.set_node(x, y, Arc::new(node)); + + width +}
diff --git a/native/vendor/datafusion-physical-plan/src/repartition/distributor_channels.rs b/native/vendor/datafusion-physical-plan/src/repartition/distributor_channels.rs new file mode 100644 index 0000000..22872d1 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/repartition/distributor_channels.rs
@@ -0,0 +1,855 @@ +// 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. + +//! Special channel construction to distribute data from various inputs into N outputs +//! minimizing buffering but preventing deadlocks when repartitioning +//! +//! # Design +//! +//! ```text +//! +----+ +------+ +//! | TX |==|| | Gate | +//! +----+ || | | +--------+ +----+ +//! ====| |==| Buffer |==| RX | +//! +----+ || | | +--------+ +----+ +//! | TX |==|| | | +//! +----+ | | +//! | | +//! +----+ | | +--------+ +----+ +//! | TX |======| |==| Buffer |==| RX | +//! +----+ +------+ +--------+ +----+ +//! ``` +//! +//! There are `N` virtual MPSC (multi-producer, single consumer) channels with unbounded capacity. However, if all +//! buffers/channels are non-empty, than a global gate will be closed preventing new data from being written (the +//! sender futures will be [pending](Poll::Pending)) until at least one channel is empty (and not closed). +use std::{ + collections::VecDeque, + future::Future, + ops::DerefMut, + pin::Pin, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + task::{Context, Poll, Waker}, +}; + +use parking_lot::Mutex; + +/// Create `n` empty channels. +pub fn channels<T>( + n: usize, +) -> (Vec<DistributionSender<T>>, Vec<DistributionReceiver<T>>) { + let channels = (0..n) + .map(|id| Arc::new(Channel::new_with_one_sender(id))) + .collect::<Vec<_>>(); + let gate = Arc::new(Gate { + empty_channels: AtomicUsize::new(n), + send_wakers: Mutex::new(None), + }); + let senders = channels + .iter() + .map(|channel| DistributionSender { + channel: Arc::clone(channel), + gate: Arc::clone(&gate), + }) + .collect(); + let receivers = channels + .into_iter() + .map(|channel| DistributionReceiver { + channel, + gate: Arc::clone(&gate), + }) + .collect(); + (senders, receivers) +} + +type PartitionAwareSenders<T> = Vec<Vec<DistributionSender<T>>>; +type PartitionAwareReceivers<T> = Vec<Vec<DistributionReceiver<T>>>; + +/// Create `n_out` empty channels for each of the `n_in` inputs. +/// This way, each distinct partition will communicate via a dedicated channel. +/// This SPSC structure enables us to track which partition input data comes from. +pub fn partition_aware_channels<T>( + n_in: usize, + n_out: usize, +) -> (PartitionAwareSenders<T>, PartitionAwareReceivers<T>) { + (0..n_in).map(|_| channels(n_out)).unzip() +} + +/// Erroring during [send](DistributionSender::send). +/// +/// This occurs when the [receiver](DistributionReceiver) is gone. +#[derive(PartialEq, Eq)] +pub struct SendError<T>(pub T); + +impl<T> std::fmt::Debug for SendError<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("SendError").finish() + } +} + +impl<T> std::fmt::Display for SendError<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "cannot send data, receiver is gone") + } +} + +impl<T> std::error::Error for SendError<T> {} + +/// Sender side of distribution [channels]. +/// +/// This handle can be cloned. All clones will write into the same channel. Dropping the last sender will close the +/// channel. In this case, the [receiver](DistributionReceiver) will still be able to poll the remaining data, but will +/// receive `None` afterwards. +#[derive(Debug)] +pub struct DistributionSender<T> { + /// To prevent lock inversion / deadlock, channel lock is always acquired prior to gate lock + channel: SharedChannel<T>, + gate: SharedGate, +} + +impl<T> DistributionSender<T> { + /// Send data. + /// + /// This fails if the [receiver](DistributionReceiver) is gone. + pub fn send(&self, element: T) -> SendFuture<'_, T> { + SendFuture { + channel: &self.channel, + gate: &self.gate, + element: Box::new(Some(element)), + } + } +} + +impl<T> Clone for DistributionSender<T> { + fn clone(&self) -> Self { + self.channel.n_senders.fetch_add(1, Ordering::SeqCst); + + Self { + channel: Arc::clone(&self.channel), + gate: Arc::clone(&self.gate), + } + } +} + +impl<T> Drop for DistributionSender<T> { + fn drop(&mut self) { + let n_senders_pre = self.channel.n_senders.fetch_sub(1, Ordering::SeqCst); + // is the last copy of the sender side? + if n_senders_pre > 1 { + return; + } + + let receivers = { + let mut state = self.channel.state.lock(); + + // During the shutdown of a empty channel, both the sender and the receiver side will be dropped. However we + // only want to decrement the "empty channels" counter once. + // + // We are within a critical section here, so we we can safely assume that either the last sender or the + // receiver (there's only one) will be dropped first. + // + // If the last sender is dropped first, `state.data` will still exists and the sender side decrements the + // signal. The receiver side then MUST check the `n_senders` counter during the section and if it is zero, + // it infers that it is dropped afterwards and MUST NOT decrement the counter. + // + // If the receiver end is dropped first, it will infer -- based on `n_senders` -- that there are still + // senders and it will decrement the `empty_channels` counter. It will also set `data` to `None`. The sender + // side will then see that `data` is `None` and can therefore infer that the receiver end was dropped, and + // hence it MUST NOT decrement the `empty_channels` counter. + if state + .data + .as_ref() + .map(|data| data.is_empty()) + .unwrap_or_default() + { + // channel is gone, so we need to clear our signal + self.gate.decr_empty_channels(); + } + + // make sure that nobody can add wakers anymore + state.recv_wakers.take().expect("not closed yet") + }; + + // wake outside of lock scope + for recv in receivers { + recv.wake(); + } + } +} + +/// Future backing [send](DistributionSender::send). +#[derive(Debug)] +pub struct SendFuture<'a, T> { + channel: &'a SharedChannel<T>, + gate: &'a SharedGate, + // the additional Box is required for `Self: Unpin` + element: Box<Option<T>>, +} + +impl<T> Future for SendFuture<'_, T> { + type Output = Result<(), SendError<T>>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + assert!(this.element.is_some(), "polled ready future"); + + // lock scope + let to_wake = { + let mut guard_channel_state = this.channel.state.lock(); + + let Some(data) = guard_channel_state.data.as_mut() else { + // receiver end dead + return Poll::Ready(Err(SendError( + this.element.take().expect("just checked"), + ))); + }; + + // does ANY receiver need data? + // if so, allow sender to create another + if this.gate.empty_channels.load(Ordering::SeqCst) == 0 { + let mut guard = this.gate.send_wakers.lock(); + if let Some(send_wakers) = guard.deref_mut() { + send_wakers.push((cx.waker().clone(), this.channel.id)); + return Poll::Pending; + } + } + + let was_empty = data.is_empty(); + data.push_back(this.element.take().expect("just checked")); + + if was_empty { + this.gate.decr_empty_channels(); + guard_channel_state.take_recv_wakers() + } else { + Vec::with_capacity(0) + } + }; + + // wake outside of lock scope + for receiver in to_wake { + receiver.wake(); + } + + Poll::Ready(Ok(())) + } +} + +/// Receiver side of distribution [channels]. +#[derive(Debug)] +pub struct DistributionReceiver<T> { + channel: SharedChannel<T>, + gate: SharedGate, +} + +impl<T> DistributionReceiver<T> { + /// Receive data from channel. + /// + /// Returns `None` if the channel is empty and no [senders](DistributionSender) are left. + pub fn recv(&mut self) -> RecvFuture<'_, T> { + RecvFuture { + channel: &mut self.channel, + gate: &mut self.gate, + rdy: false, + } + } +} + +impl<T> Drop for DistributionReceiver<T> { + fn drop(&mut self) { + let mut guard_channel_state = self.channel.state.lock(); + let data = guard_channel_state.data.take().expect("not dropped yet"); + + // See `DistributedSender::drop` for an explanation of the drop order and when the "empty channels" counter is + // decremented. + if data.is_empty() && (self.channel.n_senders.load(Ordering::SeqCst) > 0) { + // channel is gone, so we need to clear our signal + self.gate.decr_empty_channels(); + } + + // senders may be waiting for gate to open but should error now that the channel is closed + self.gate.wake_channel_senders(self.channel.id); + } +} + +/// Future backing [recv](DistributionReceiver::recv). +pub struct RecvFuture<'a, T> { + channel: &'a mut SharedChannel<T>, + gate: &'a mut SharedGate, + rdy: bool, +} + +impl<T> Future for RecvFuture<'_, T> { + type Output = Option<T>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = &mut *self; + assert!(!this.rdy, "polled ready future"); + + let mut guard_channel_state = this.channel.state.lock(); + let channel_state = guard_channel_state.deref_mut(); + let data = channel_state.data.as_mut().expect("not dropped yet"); + + match data.pop_front() { + Some(element) => { + // change "empty" signal for this channel? + if data.is_empty() && channel_state.recv_wakers.is_some() { + // update counter + let old_counter = + this.gate.empty_channels.fetch_add(1, Ordering::SeqCst); + + // open gate? + let to_wake = if old_counter == 0 { + let mut guard = this.gate.send_wakers.lock(); + + // check after lock to see if we should still change the state + if this.gate.empty_channels.load(Ordering::SeqCst) > 0 { + guard.take().unwrap_or_default() + } else { + Vec::with_capacity(0) + } + } else { + Vec::with_capacity(0) + }; + + drop(guard_channel_state); + + // wake outside of lock scope + for (waker, _channel_id) in to_wake { + waker.wake(); + } + } + + this.rdy = true; + Poll::Ready(Some(element)) + } + None => { + if let Some(recv_wakers) = channel_state.recv_wakers.as_mut() { + recv_wakers.push(cx.waker().clone()); + Poll::Pending + } else { + this.rdy = true; + Poll::Ready(None) + } + } + } + } +} + +/// Links senders and receivers. +#[derive(Debug)] +struct Channel<T> { + /// Reference counter for the sender side. + n_senders: AtomicUsize, + + /// Channel ID. + /// + /// This is used to address [send wakers](Gate::send_wakers). + id: usize, + + /// Mutable state. + state: Mutex<ChannelState<T>>, +} + +impl<T> Channel<T> { + /// Create new channel with one sender (so we don't need to [fetch-add](AtomicUsize::fetch_add) directly afterwards). + fn new_with_one_sender(id: usize) -> Self { + Channel { + n_senders: AtomicUsize::new(1), + id, + state: Mutex::new(ChannelState { + data: Some(VecDeque::default()), + recv_wakers: Some(Vec::default()), + }), + } + } +} + +#[derive(Debug)] +struct ChannelState<T> { + /// Buffered data. + /// + /// This is [`None`] when the receiver is gone. + data: Option<VecDeque<T>>, + + /// Wakers for the receiver side. + /// + /// The receiver will be pending if the [buffer](Self::data) is empty and + /// there are senders left (otherwise this is set to [`None`]). + recv_wakers: Option<Vec<Waker>>, +} + +impl<T> ChannelState<T> { + /// Get all [`recv_wakers`](Self::recv_wakers) and replace with identically-sized buffer. + /// + /// The wakers should be woken AFTER the lock to [this state](Self) was dropped. + /// + /// # Panics + /// Assumes that channel is NOT closed yet, i.e. that [`recv_wakers`](Self::recv_wakers) is not [`None`]. + fn take_recv_wakers(&mut self) -> Vec<Waker> { + let to_wake = self.recv_wakers.as_mut().expect("not closed"); + let mut tmp = Vec::with_capacity(to_wake.capacity()); + std::mem::swap(to_wake, &mut tmp); + tmp + } +} + +/// Shared channel. +/// +/// One or multiple senders and a single receiver will share a channel. +type SharedChannel<T> = Arc<Channel<T>>; + +/// The "all channels have data" gate. +#[derive(Debug)] +struct Gate { + /// Number of currently empty (and still open) channels. + empty_channels: AtomicUsize, + + /// Wakers for the sender side, including their channel IDs. + /// + /// This is `None` if the there are non-empty channels. + send_wakers: Mutex<Option<Vec<(Waker, usize)>>>, +} + +impl Gate { + /// Wake senders for a specific channel. + /// + /// This is helpful to signal that the receiver side is gone and the senders shall now error. + fn wake_channel_senders(&self, id: usize) { + // lock scope + let to_wake = { + let mut guard = self.send_wakers.lock(); + + if let Some(send_wakers) = guard.deref_mut() { + // `drain_filter` is unstable, so implement our own + let (wake, keep) = + send_wakers.drain(..).partition(|(_waker, id2)| id == *id2); + + *send_wakers = keep; + + wake + } else { + Vec::with_capacity(0) + } + }; + + // wake outside of lock scope + for (waker, _id) in to_wake { + waker.wake(); + } + } + + fn decr_empty_channels(&self) { + let old_count = self.empty_channels.fetch_sub(1, Ordering::SeqCst); + + if old_count == 1 { + let mut guard = self.send_wakers.lock(); + + // double-check state during lock + if self.empty_channels.load(Ordering::SeqCst) == 0 && guard.is_none() { + *guard = Some(Vec::new()); + } + } + } +} + +/// Gate shared by all senders and receivers. +type SharedGate = Arc<Gate>; + +#[cfg(test)] +mod tests { + use std::sync::atomic::AtomicBool; + + use futures::{FutureExt, task::ArcWake}; + + use super::*; + + #[test] + fn test_single_channel_no_gate() { + // use two channels so that the first one never hits the gate + let (mut txs, mut rxs) = channels(2); + + let mut recv_fut = rxs[0].recv(); + let waker = poll_pending(&mut recv_fut); + + poll_ready(&mut txs[0].send("foo")).unwrap(); + assert!(waker.woken()); + assert_eq!(poll_ready(&mut recv_fut), Some("foo"),); + + poll_ready(&mut txs[0].send("bar")).unwrap(); + poll_ready(&mut txs[0].send("baz")).unwrap(); + poll_ready(&mut txs[0].send("end")).unwrap(); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("bar"),); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("baz"),); + + // close channel + txs.remove(0); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("end"),); + assert_eq!(poll_ready(&mut rxs[0].recv()), None,); + assert_eq!(poll_ready(&mut rxs[0].recv()), None,); + } + + #[test] + fn test_multi_sender() { + // use two channels so that the first one never hits the gate + let (txs, mut rxs) = channels(2); + + let tx_clone = txs[0].clone(); + + poll_ready(&mut txs[0].send("foo")).unwrap(); + poll_ready(&mut tx_clone.send("bar")).unwrap(); + + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("foo"),); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("bar"),); + } + + #[test] + fn test_gate() { + let (txs, mut rxs) = channels(2); + + // gate initially open + poll_ready(&mut txs[0].send("0_a")).unwrap(); + + // gate still open because channel 1 is still empty + poll_ready(&mut txs[0].send("0_b")).unwrap(); + + // gate still open because channel 1 is still empty prior to this call, so this call still goes through + poll_ready(&mut txs[1].send("1_a")).unwrap(); + + // both channels non-empty => gate closed + + let mut send_fut = txs[1].send("1_b"); + let waker = poll_pending(&mut send_fut); + + // drain channel 0 + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("0_a"),); + poll_pending(&mut send_fut); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("0_b"),); + + // channel 0 empty => gate open + assert!(waker.woken()); + poll_ready(&mut send_fut).unwrap(); + } + + #[test] + fn test_close_channel_by_dropping_tx() { + let (mut txs, mut rxs) = channels(2); + + let tx0 = txs.remove(0); + let tx1 = txs.remove(0); + let tx0_clone = tx0.clone(); + + let mut recv_fut = rxs[0].recv(); + + poll_ready(&mut tx1.send("a")).unwrap(); + let recv_waker = poll_pending(&mut recv_fut); + + // drop original sender + drop(tx0); + + // not yet closed (there's a clone left) + assert!(!recv_waker.woken()); + poll_ready(&mut tx1.send("b")).unwrap(); + let recv_waker = poll_pending(&mut recv_fut); + + // create new clone + let tx0_clone2 = tx0_clone.clone(); + assert!(!recv_waker.woken()); + poll_ready(&mut tx1.send("c")).unwrap(); + let recv_waker = poll_pending(&mut recv_fut); + + // drop first clone + drop(tx0_clone); + assert!(!recv_waker.woken()); + poll_ready(&mut tx1.send("d")).unwrap(); + let recv_waker = poll_pending(&mut recv_fut); + + // drop last clone + drop(tx0_clone2); + + // channel closed => also close gate + poll_pending(&mut tx1.send("e")); + assert!(recv_waker.woken()); + assert_eq!(poll_ready(&mut recv_fut), None,); + } + + #[test] + fn test_close_channel_by_dropping_rx_on_open_gate() { + let (txs, mut rxs) = channels(2); + + let rx0 = rxs.remove(0); + let _rx1 = rxs.remove(0); + + poll_ready(&mut txs[1].send("a")).unwrap(); + + // drop receiver => also close gate + drop(rx0); + + poll_pending(&mut txs[1].send("b")); + assert_eq!(poll_ready(&mut txs[0].send("foo")), Err(SendError("foo")),); + } + + #[test] + fn test_close_channel_by_dropping_rx_on_closed_gate() { + let (txs, mut rxs) = channels(2); + + let rx0 = rxs.remove(0); + let mut rx1 = rxs.remove(0); + + // fill both channels + poll_ready(&mut txs[0].send("0_a")).unwrap(); + poll_ready(&mut txs[1].send("1_a")).unwrap(); + + let mut send_fut0 = txs[0].send("0_b"); + let mut send_fut1 = txs[1].send("1_b"); + let waker0 = poll_pending(&mut send_fut0); + let waker1 = poll_pending(&mut send_fut1); + + // drop receiver + drop(rx0); + + assert!(waker0.woken()); + assert!(!waker1.woken()); + assert_eq!(poll_ready(&mut send_fut0), Err(SendError("0_b")),); + + // gate closed, so cannot send on channel 1 + poll_pending(&mut send_fut1); + + // channel 1 can still receive data + assert_eq!(poll_ready(&mut rx1.recv()), Some("1_a"),); + } + + #[test] + fn test_drop_rx_three_channels() { + let (mut txs, mut rxs) = channels(3); + + let tx0 = txs.remove(0); + let tx1 = txs.remove(0); + let tx2 = txs.remove(0); + let mut rx0 = rxs.remove(0); + let rx1 = rxs.remove(0); + let _rx2 = rxs.remove(0); + + // fill channels + poll_ready(&mut tx0.send("0_a")).unwrap(); + poll_ready(&mut tx1.send("1_a")).unwrap(); + poll_ready(&mut tx2.send("2_a")).unwrap(); + + // drop / close one channel + drop(rx1); + + // receive data + assert_eq!(poll_ready(&mut rx0.recv()), Some("0_a"),); + + // use senders again + poll_ready(&mut tx0.send("0_b")).unwrap(); + assert_eq!(poll_ready(&mut tx1.send("1_b")), Err(SendError("1_b")),); + poll_pending(&mut tx2.send("2_b")); + } + + #[test] + fn test_close_channel_by_dropping_rx_clears_data() { + let (txs, rxs) = channels(1); + + let obj = Arc::new(()); + let counter = Arc::downgrade(&obj); + assert_eq!(counter.strong_count(), 1); + + // add object to channel + poll_ready(&mut txs[0].send(obj)).unwrap(); + assert_eq!(counter.strong_count(), 1); + + // drop receiver + drop(rxs); + + assert_eq!(counter.strong_count(), 0); + } + + /// Ensure that polling "pending" futures work even when you poll them too often (which happens under some circumstances). + #[test] + fn test_poll_empty_channel_twice() { + let (txs, mut rxs) = channels(1); + + let mut recv_fut = rxs[0].recv(); + let waker_1a = poll_pending(&mut recv_fut); + let waker_1b = poll_pending(&mut recv_fut); + + let mut recv_fut = rxs[0].recv(); + let waker_2 = poll_pending(&mut recv_fut); + + poll_ready(&mut txs[0].send("a")).unwrap(); + assert!(waker_1a.woken()); + assert!(waker_1b.woken()); + assert!(waker_2.woken()); + assert_eq!(poll_ready(&mut recv_fut), Some("a"),); + + poll_ready(&mut txs[0].send("b")).unwrap(); + let mut send_fut = txs[0].send("c"); + let waker_3 = poll_pending(&mut send_fut); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("b"),); + assert!(waker_3.woken()); + poll_ready(&mut send_fut).unwrap(); + assert_eq!(poll_ready(&mut rxs[0].recv()), Some("c")); + + let mut recv_fut = rxs[0].recv(); + let waker_4 = poll_pending(&mut recv_fut); + + let mut recv_fut = rxs[0].recv(); + let waker_5 = poll_pending(&mut recv_fut); + + poll_ready(&mut txs[0].send("d")).unwrap(); + let mut send_fut = txs[0].send("e"); + let waker_6a = poll_pending(&mut send_fut); + let waker_6b = poll_pending(&mut send_fut); + + assert!(waker_4.woken()); + assert!(waker_5.woken()); + assert_eq!(poll_ready(&mut recv_fut), Some("d"),); + + assert!(waker_6a.woken()); + assert!(waker_6b.woken()); + poll_ready(&mut send_fut).unwrap(); + } + + #[test] + #[should_panic(expected = "polled ready future")] + fn test_panic_poll_send_future_after_ready_ok() { + let (txs, _rxs) = channels(1); + let mut fut = txs[0].send("foo"); + poll_ready(&mut fut).unwrap(); + poll_ready(&mut fut).ok(); + } + + #[test] + #[should_panic(expected = "polled ready future")] + fn test_panic_poll_send_future_after_ready_err() { + let (txs, rxs) = channels(1); + + drop(rxs); + + let mut fut = txs[0].send("foo"); + poll_ready(&mut fut).unwrap_err(); + poll_ready(&mut fut).ok(); + } + + #[test] + #[should_panic(expected = "polled ready future")] + fn test_panic_poll_recv_future_after_ready_some() { + let (txs, mut rxs) = channels(1); + + poll_ready(&mut txs[0].send("foo")).unwrap(); + + let mut fut = rxs[0].recv(); + poll_ready(&mut fut).unwrap(); + poll_ready(&mut fut); + } + + #[test] + #[should_panic(expected = "polled ready future")] + fn test_panic_poll_recv_future_after_ready_none() { + let (txs, mut rxs) = channels::<u8>(1); + + drop(txs); + + let mut fut = rxs[0].recv(); + assert!(poll_ready(&mut fut).is_none()); + poll_ready(&mut fut); + } + + #[test] + #[should_panic(expected = "future is pending")] + fn test_meta_poll_ready_wrong_state() { + let mut fut = futures::future::pending::<u8>(); + poll_ready(&mut fut); + } + + #[test] + #[should_panic(expected = "future is ready")] + fn test_meta_poll_pending_wrong_state() { + let mut fut = futures::future::ready(1); + poll_pending(&mut fut); + } + + /// Test [`poll_pending`] (i.e. the testing utils, not the actual library code). + #[test] + fn test_meta_poll_pending_waker() { + let (tx, mut rx) = futures::channel::oneshot::channel(); + let waker = poll_pending(&mut rx); + assert!(!waker.woken()); + tx.send(1).unwrap(); + assert!(waker.woken()); + } + + /// Poll a given [`Future`] and ensure it is [ready](Poll::Ready). + #[track_caller] + fn poll_ready<F>(fut: &mut F) -> F::Output + where + F: Future + Unpin, + { + match poll(fut).0 { + Poll::Ready(x) => x, + Poll::Pending => panic!("future is pending"), + } + } + + /// Poll a given [`Future`] and ensure it is [pending](Poll::Pending). + /// + /// Returns a waker that can later be checked. + #[track_caller] + fn poll_pending<F>(fut: &mut F) -> Arc<TestWaker> + where + F: Future + Unpin, + { + let (res, waker) = poll(fut); + match res { + Poll::Ready(_) => panic!("future is ready"), + Poll::Pending => waker, + } + } + + fn poll<F>(fut: &mut F) -> (Poll<F::Output>, Arc<TestWaker>) + where + F: Future + Unpin, + { + let test_waker = Arc::new(TestWaker::default()); + let waker = futures::task::waker(Arc::clone(&test_waker)); + let mut cx = Context::from_waker(&waker); + let res = fut.poll_unpin(&mut cx); + (res, test_waker) + } + + /// A test [`Waker`] that signal if [`wake`](Waker::wake) was called. + #[derive(Debug, Default)] + struct TestWaker { + woken: AtomicBool, + } + + impl TestWaker { + /// Was [`wake`](Waker::wake) called? + fn woken(&self) -> bool { + self.woken.load(Ordering::SeqCst) + } + } + + impl ArcWake for TestWaker { + fn wake_by_ref(arc_self: &Arc<Self>) { + arc_self.woken.store(true, Ordering::SeqCst); + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/repartition/mod.rs b/native/vendor/datafusion-physical-plan/src/repartition/mod.rs new file mode 100644 index 0000000..063954a --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/repartition/mod.rs
@@ -0,0 +1,4612 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This file implements the [`RepartitionExec`] operator, which maps N input +//! partitions to M output partitions based on a partitioning scheme, optionally +//! maintaining the order of the input rows in the output. + +use std::cmp::Ordering; +use std::fmt::{Debug, Display, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use std::task::{Context, Poll}; +use std::vec; + +use super::common::SharedMemoryReservation; +use super::metrics::{self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet}; +use super::{ + DisplayAs, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream, +}; +use crate::coalesce::LimitedBatchCoalescer; +use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::hash_utils::create_hashes; +use crate::metrics::{BaselineMetrics, SpillMetrics}; +use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; +use crate::sorts::streaming_merge::StreamingMergeBuilder; +use crate::spill::spill_manager::SpillManager; +use crate::spill::spill_pool::{self, SpillPoolSink, SpillPoolWriter}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count, +}; + +use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; +use arrow::compute::take_arrays; +use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type}; +use arrow_schema::SortOptions; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; +use datafusion_common::{ + ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, + assert_or_internal_err, internal_datafusion_err, internal_err, + validate_range_split_points, +}; +use datafusion_common::{Result, not_impl_err}; +use datafusion_common_runtime::SpawnedTask; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_expr::ColumnarValue; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr, RangePartitioning}; +use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf; + +use crate::filter_pushdown::{ + ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, +}; +use crate::joins::SeededRandomState; +use crate::sort_pushdown::SortOrderPushdownResult; +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::stream::Stream; +use futures::{FutureExt, StreamExt, TryStreamExt}; +use log::trace; +use parking_lot::Mutex; + +mod distributor_channels; +use crate::repartition::distributor_channels::SendError; +use distributor_channels::{ + DistributionReceiver, DistributionSender, channels, partition_aware_channels, +}; + +/// A batch in the repartition queue - either in memory or spilled to disk. +/// +/// This enum represents the two states a batch can be in during repartitioning. +/// The decision to spill is made based on memory availability when sending a batch +/// to an output partition. +/// +/// # Batch Flow with Spilling +/// +/// ```text +/// Input Stream ◀──────┐ +/// │ │ +/// ▼ │ +/// Partition Logic │ +/// │ `batch_size` not +/// ▼ reached yet +/// Coalesce Batch │ +/// ┌───────────────┴────────────────┘ +/// ▼ +/// `batch_size` reached +/// │ +/// └───────────────┐ +/// ▼ +/// try_grow() +/// ┌───────────────┴────────────────┐ +/// ▼ ▼ +/// try_grow() succeeds try_grow() fails +/// (Memory Available) (Memory Pressure) +/// │ │ +/// ▼ ▼ +/// RepartitionBatch::Memory spill_writer.push_batch() +/// (batch held in memory) (batch written to disk) +/// │ │ +/// │ ▼ +/// │ RepartitionBatch::Spilled +/// │ (marker - no batch data) +/// └──────────────┬─────────────────┘ +/// │ +/// ▼ +/// Send to channel +/// │ +/// ▼ +/// Output Stream (poll) +/// │ +/// ┌──────────────┴────────────────┐ +/// ▼ ▼ +/// RepartitionBatch::Memory RepartitionBatch::Spilled +/// Return batch immediately Poll spill_stream (blocks) +/// └─────────────┬─────────────────┘ +/// │ +/// ▼ +/// Return batch +/// (FIFO order preserved) +/// ``` +/// +/// See [`RepartitionExec`] for overall architecture and [`StreamState`] for +/// the state machine that handles reading these batches. +#[derive(Debug)] +enum RepartitionBatch { + /// Batch held in memory (counts against memory reservation) + Memory(RecordBatch), + /// Marker indicating a batch was spilled to the partition's SpillPool. + /// The actual batch can be retrieved by reading from the SpillPoolStream. + /// This variant contains no data itself - it's just a signal to the reader + /// to fetch the next batch from the spill stream. + Spilled, +} + +type MaybeBatch = Option<Result<RepartitionBatch>>; +type InputPartitionsToCurrentPartitionSender = Vec<DistributionSender<MaybeBatch>>; +type InputPartitionsToCurrentPartitionReceiver = Vec<DistributionReceiver<MaybeBatch>>; + +/// Output channel with its associated memory reservation and spill writer. +/// +/// `coalescer` is `None` for preserve-order mode, where downstream +/// [`StreamingMergeBuilder`] performs the batching; otherwise it's a +/// [`SharedCoalescer`] cloned from the per-partition one held by +/// [`PartitionChannels`]. +struct OutputChannel { + sender: DistributionSender<MaybeBatch>, + reservation: SharedMemoryReservation, + spill_writer: SpillPoolSink, + shared_coalescer: Option<SharedCoalescer>, +} + +/// The set of spill-pool writers for a single output partition, before they are handed to the +/// per-input tasks. The variant encodes the repartition mode so the wrong writer topology cannot +/// be constructed for a given mode. +enum PartitionSpillWriters { + /// `preserve_order`: one single-producer FIFO writer per input partition. Each is `take`n + /// exactly once (moved into the matching input task), so the pool always has one writer. + PerInput(Vec<Option<SpillPoolSink>>), + /// Non-preserve-order: one shared writer, cloned into every input task. + Shared(SpillPoolWriter), +} + +impl PartitionSpillWriters { + /// Hand out the writer for input partition `input`. + /// + /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per + /// input); in `Shared` mode it clones the shared writer. + fn take_for_input(&mut self, input: usize) -> Result<SpillPoolSink> { + match self { + PartitionSpillWriters::PerInput(writers) => { + writers[input].take().ok_or_else(|| { + internal_datafusion_err!( + "spill writer for input partition requested more than once" + ) + }) + } + PartitionSpillWriters::Shared(writer) => Ok(writer.new_sink()), + } + } +} + +impl OutputChannel { + fn coalesce(&mut self, batch: RecordBatch) -> Result<Vec<RecordBatch>> { + match &self.shared_coalescer { + Some(shared) => Ok(shared.push_and_drain(batch)?), + None => Ok(vec![batch]), + } + } + + /// Send a single batch through the channel for `partition`, applying + /// the memory reservation / spill-writer fallback. Removes the channel + /// from `self.inner` if the receiver has hung up. + /// + /// Used after [`OutputChannel::coalesce`] for performance purposes. + async fn send(&mut self, batch: RecordBatch) -> Result<(), SendError<MaybeBatch>> { + let size = batch.get_array_memory_size(); + + // Decide the payload outside of any await: never hold a MutexGuard + // across an await point. + let (payload, is_memory_batch) = { + match self.reservation.try_grow(size) { + Ok(_) => (Ok(RepartitionBatch::Memory(batch)), true), + Err(_) => match self.spill_writer.push_batch(&batch) { + Ok(()) => (Ok(RepartitionBatch::Spilled), false), + Err(err) => (Err(err), false), + }, + } + }; + + let result = self.sender.send(Some(payload)).await; + if result.is_err() && is_memory_batch { + self.reservation.shrink(size); + } + result + } + + async fn finalize(mut self) -> Result<()> { + let Some(shared) = self.shared_coalescer.take() else { + return Ok(()); + }; + for batch in shared.finalize()? { + // If this errored, it means that nobody is listening on the other side, which is fine + // and can happen in certain cases, like when a LIMIT drops the stream that listens. + let _ = self.send(batch).await; + } + Ok(()) + } +} + +/// A producer-side coalescer shared across all input tasks targeting a +/// single output partition. +/// +/// Bundles the [`LimitedBatchCoalescer`] (behind a [`Mutex`]) with the +/// active-sender counter that tracks how many input tasks may still push +/// into it. The last task to call [`Self::finalize`] is the one that +/// finalizes the coalescer and ships the residual batch. +/// +/// Cheap to [`Clone`]: both fields are [`Arc`]s. +#[derive(Clone)] +struct SharedCoalescer { + inner: Arc<Mutex<LimitedBatchCoalescer>>, + active_senders: Arc<AtomicUsize>, +} + +impl SharedCoalescer { + fn new(schema: SchemaRef, target_batch_size: usize, num_senders: usize) -> Self { + Self { + inner: Arc::new(Mutex::new(LimitedBatchCoalescer::new( + schema, + target_batch_size, + None, + ))), + active_senders: Arc::new(AtomicUsize::new(num_senders)), + } + } + + /// Push `batch` into the coalescer and drain any newly completed + /// batches. The mutex is held only briefly. + fn push_and_drain(&self, batch: RecordBatch) -> Result<Vec<RecordBatch>> { + let mut acc = Vec::new(); + let mut c = self.inner.lock(); + c.push_batch(batch)?; + while let Some(b) = c.next_completed_batch() { + acc.push(b); + } + Ok(acc) + } + + /// Decrement the active-senders counter. If this caller was the last + /// sender, finalize the coalescer and return its residual batches; if + /// other senders are still active, return `Ok(None)`. + fn finalize(&self) -> Result<Vec<RecordBatch>> { + let was_last = self.active_senders.fetch_sub(1, AtomicOrdering::AcqRel) == 1; + if !was_last { + return Ok(vec![]); + } + let mut acc = Vec::new(); + let mut c = self.inner.lock(); + c.finish()?; + while let Some(b) = c.next_completed_batch() { + acc.push(b); + } + Ok(acc) + } +} + +/// Channels and resources for a single output partition. +/// +/// Each output partition has channels to receive data from all input partitions. +/// To handle memory pressure, each (input, output) pair gets its own +/// [`SpillPool`](crate::spill::spill_pool) channel via [`spill_pool::channel`]. +/// +/// # Structure +/// +/// For an output partition receiving from N input partitions: +/// - `tx`: N senders (one per input) for sending batches to this output +/// - `rx`: N receivers (one per input) for receiving batches at this output +/// - `spill_writers`: N spill writers (one per input) for writing spilled data +/// - `spill_readers`: N spill readers (one per input) for reading spilled data +/// +/// This 1:1 mapping between input partitions and spill channels ensures that +/// batches from each input are processed in FIFO order, even when some batches +/// are spilled to disk and others remain in memory. +/// +/// See [`RepartitionExec`] for the overall N×M architecture. +/// +/// [`spill_pool::channel`]: crate::spill::spill_pool::spsc_channel +struct PartitionChannels { + /// Senders for each input partition to send data to this output partition + tx: InputPartitionsToCurrentPartitionSender, + /// Receivers for each input partition sending data to this output partition + rx: InputPartitionsToCurrentPartitionReceiver, + /// Memory reservation for this output partition + reservation: SharedMemoryReservation, + /// Shared coalescer used by all input tasks targeting this output + /// partition. `None` in preserve-order mode (downstream + /// `StreamingMergeBuilder` handles batching). + shared_coalescer: Option<SharedCoalescer>, + /// Spill writers for writing spilled data, before they are handed to the per-input tasks. + /// The variant is chosen by the repartition mode (see [`PartitionSpillWriters`]): a dedicated + /// single-producer FIFO writer per input in preserve-order mode, or one shared writer in + /// non-preserve-order mode. + spill_writers: PartitionSpillWriters, + /// Spill readers for reading spilled data - one per input partition (FIFO semantics). + /// Each (input, output) pair gets its own reader to maintain proper ordering. + spill_readers: Vec<SendableRecordBatchStream>, +} + +struct ConsumingInputStreamsState { + /// Channels for sending batches from input partitions to output partitions. + /// Key is the partition number. + channels: HashMap<usize, PartitionChannels>, + + /// Helper that ensures that background jobs are killed once they are no longer needed. + abort_helper: Arc<Vec<SpawnedTask<()>>>, +} + +impl Debug for ConsumingInputStreamsState { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ConsumingInputStreamsState") + .field("num_channels", &self.channels.len()) + .field("abort_helper", &self.abort_helper) + .finish() + } +} + +/// Inner state of [`RepartitionExec`]. +#[derive(Default)] +enum RepartitionExecState { + /// Not initialized yet. This is the default state stored in the RepartitionExec node + /// upon instantiation. + #[default] + NotInitialized, + /// Input streams are initialized, but they are still not being consumed. The node + /// transitions to this state when the arrow's RecordBatch stream is created in + /// RepartitionExec::execute(), but before any message is polled. + InputStreamsInitialized(Vec<(SendableRecordBatchStream, RepartitionMetrics)>), + /// The input streams are being consumed. The node transitions to this state when + /// the first message in the arrow's RecordBatch stream is consumed. + ConsumingInputStreams(ConsumingInputStreamsState), +} + +impl Debug for RepartitionExecState { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + RepartitionExecState::NotInitialized => write!(f, "NotInitialized"), + RepartitionExecState::InputStreamsInitialized(v) => { + write!(f, "InputStreamsInitialized({:?})", v.len()) + } + RepartitionExecState::ConsumingInputStreams(v) => { + write!(f, "ConsumingInputStreams({v:?})") + } + } + } +} + +impl RepartitionExecState { + fn ensure_input_streams_initialized( + &mut self, + input: &Arc<dyn ExecutionPlan>, + metrics: &ExecutionPlanMetricsSet, + output_partitions: usize, + ctx: &Arc<TaskContext>, + ) -> Result<()> { + if !matches!(self, RepartitionExecState::NotInitialized) { + return Ok(()); + } + + let num_input_partitions = input.output_partitioning().partition_count(); + let mut streams_and_metrics = Vec::with_capacity(num_input_partitions); + + for i in 0..num_input_partitions { + let metrics = RepartitionMetrics::new(i, output_partitions, metrics); + + let timer = metrics.fetch_time.timer(); + let stream = input.execute(i, Arc::clone(ctx))?; + timer.done(); + + streams_and_metrics.push((stream, metrics)); + } + *self = RepartitionExecState::InputStreamsInitialized(streams_and_metrics); + Ok(()) + } + + #[expect(clippy::too_many_arguments)] + fn consume_input_streams( + &mut self, + input: &Arc<dyn ExecutionPlan>, + metrics: &ExecutionPlanMetricsSet, + partitioning: &Partitioning, + preserve_order: bool, + name: &str, + context: &Arc<TaskContext>, + spill_manager: SpillManager, + ) -> Result<&mut ConsumingInputStreamsState> { + let streams_and_metrics = match self { + RepartitionExecState::NotInitialized => { + self.ensure_input_streams_initialized( + input, + metrics, + partitioning.partition_count(), + context, + )?; + let RepartitionExecState::InputStreamsInitialized(value) = self else { + // This cannot happen, as ensure_input_streams_initialized() was just called, + // but the compiler does not know. + return internal_err!( + "Programming error: RepartitionExecState must be in the InputStreamsInitialized state after calling RepartitionExecState::ensure_input_streams_initialized" + ); + }; + value + } + RepartitionExecState::ConsumingInputStreams(value) => return Ok(value), + RepartitionExecState::InputStreamsInitialized(value) => value, + }; + + let num_input_partitions = streams_and_metrics.len(); + let num_output_partitions = partitioning.partition_count(); + let coalesce_batches = !preserve_order && !input.boundedness().is_unbounded(); + + let spill_manager = Arc::new(spill_manager); + + let (txs, rxs) = if preserve_order { + // Create partition-aware channels with one channel per (input, output) pair + // This provides backpressure while maintaining proper ordering + let (txs_all, rxs_all) = + partition_aware_channels(num_input_partitions, num_output_partitions); + // Take transpose of senders and receivers. `state.channels` keeps track of entries per output partition + let txs = transpose(txs_all); + let rxs = transpose(rxs_all); + (txs, rxs) + } else { + // Create one channel per *output* partition with backpressure + let (txs, rxs) = channels(num_output_partitions); + // Clone sender for each input partitions + let txs = txs + .into_iter() + .map(|item| vec![item; num_input_partitions]) + .collect::<Vec<_>>(); + let rxs = rxs.into_iter().map(|item| vec![item]).collect::<Vec<_>>(); + (txs, rxs) + }; + + let mut channels = HashMap::with_capacity(txs.len()); + for (partition, (tx, rx)) in txs.into_iter().zip(rxs).enumerate() { + let reservation = Arc::new( + MemoryConsumer::new(format!("{name}[{partition}]")) + .with_can_spill(true) + .register(context.memory_pool()), + ); + + // Create spill channels based on mode: + // - preserve_order: one spill channel per (input, output) pair for proper FIFO ordering + // - non-preserve-order: one shared spill channel per output partition since all inputs + // share the same receiver + let max_file_size = context + .session_config() + .options() + .execution + .max_spill_file_size_bytes + .get(); + + let (spill_writers, spill_readers) = if preserve_order { + // preserve_order: one dedicated single-producer FIFO pool per input partition. + // Each writer is moved into exactly one input task (never cloned), so the ordering + // the downstream merge relies on is preserved across the spill boundary. + let mut writers = Vec::with_capacity(num_input_partitions); + let mut readers = Vec::with_capacity(num_input_partitions); + for _ in 0..num_input_partitions { + let (writer, reader) = spill_pool::spsc_channel( + max_file_size, + Arc::clone(&spill_manager), + ); + writers.push(Some(writer)); + readers.push(reader); + } + (PartitionSpillWriters::PerInput(writers), readers) + } else { + // non-preserve-order: one shared multi-producer pool per output partition, since + // all inputs share the same receiver and the output is an unordered multiset. + let (writer, reader) = + spill_pool::mpsc_channel(max_file_size, Arc::clone(&spill_manager)); + (PartitionSpillWriters::Shared(writer), vec![reader]) + }; + + // Coalesce on the producer side, before the channel's gate, so + // the consumer never sees the per-input-task small batches. + // Skip in preserve-order mode, where `StreamingMergeBuilder` + // handles batching, and for unbounded inputs, where a residual + // batch could otherwise be withheld indefinitely. + let shared_coalescer = coalesce_batches.then(|| { + SharedCoalescer::new( + input.schema(), + context.session_config().batch_size(), + num_input_partitions, + ) + }); + + channels.insert( + partition, + PartitionChannels { + tx, + rx, + reservation, + spill_readers, + spill_writers, + shared_coalescer, + }, + ); + } + + // launch one async task per *input* partition + let mut spawned_tasks = Vec::with_capacity(num_input_partitions); + for (i, (stream, metrics)) in + std::mem::take(streams_and_metrics).into_iter().enumerate() + { + let txs: HashMap<_, _> = channels + .iter_mut() + .map(|(partition, channels)| { + // Hand this input task its spill writer: in preserve_order mode this moves + // the input's dedicated FIFO writer out; otherwise it clones the shared + // writer. See [`PartitionSpillWriters::take_for_input`]. + Ok(( + *partition, + OutputChannel { + sender: channels.tx[i].clone(), + reservation: Arc::clone(&channels.reservation), + spill_writer: channels.spill_writers.take_for_input(i)?, + shared_coalescer: channels.shared_coalescer.clone(), + }, + )) + }) + .collect::<Result<HashMap<_, _>>>()?; + + // Extract senders for wait_for_task before moving txs + let senders: HashMap<_, _> = txs + .iter() + .map(|(partition, channel)| (*partition, channel.sender.clone())) + .collect(); + + let input_task = SpawnedTask::spawn(RepartitionExec::pull_from_input( + stream, + txs, + partitioning.clone(), + metrics, + // preserve_order depends on partition index to start from 0 + if preserve_order { 0 } else { i }, + num_input_partitions, + )); + + // In a separate task, wait for each input to be done + // (and pass along any errors, including panic!s) + let wait_for_task = + SpawnedTask::spawn(RepartitionExec::wait_for_task(input_task, senders)); + spawned_tasks.push(wait_for_task); + } + *self = Self::ConsumingInputStreams(ConsumingInputStreamsState { + channels, + abort_helper: Arc::new(spawned_tasks), + }); + match self { + RepartitionExecState::ConsumingInputStreams(value) => Ok(value), + _ => unreachable!(), + } + } +} + +/// A utility that can be used to partition batches based on [`Partitioning`] +pub struct BatchPartitioner { + state: BatchPartitionerState, + timer: metrics::Time, +} + +enum BatchPartitionerState { + Hash { + exprs: Vec<Arc<dyn PhysicalExpr>>, + partition_reducer: StrengthReducedU64, + hash_buffer: Vec<u64>, + indices: Vec<Vec<u32>>, + }, + RoundRobin { + num_partitions: usize, + next_idx: usize, + }, + Range { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Sort options from the `LexOrdering` + sort_options: Vec<SortOptions>, + /// Boundaries between adjacent partitions. + split_points: Vec<SplitPoint>, + /// Row indices grouped by output partition + indices: Vec<Vec<u32>>, + /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points + partition_buffer: Vec<ScalarValue>, + }, +} + +/// Fixed RandomState used for hash repartitioning to ensure consistent behavior across +/// executions and runs. +pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0); + +/// Physical expression that returns the Range partition for each input row. +/// +/// This uses the same routing function as [`BatchPartitioner`], so dynamic +/// filtering and repartitioning agree for every [`ScalarValue`] comparison. +#[derive(Debug, Hash, PartialEq, Eq)] +pub struct RangeExpr { + on_columns: Vec<PhysicalExprRef>, + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, +} + +impl RangeExpr { + /// Creates a Range expression for `on_columns` using the supplied routing + /// metadata. + pub fn try_new( + on_columns: Vec<PhysicalExprRef>, + range_partitioning: &RangePartitioning, + ) -> Result<Self> { + let sort_options = range_partitioning + .ordering() + .iter() + .map(|expr| expr.options) + .collect(); + Self::try_new_parts( + on_columns, + range_partitioning.split_points().to_vec(), + sort_options, + ) + } + + fn try_new_parts( + on_columns: Vec<PhysicalExprRef>, + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + ) -> Result<Self> { + assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key"); + assert_or_internal_err!( + on_columns.len() == sort_options.len(), + "RangeExpr key count must match sort options" + ); + validate_range_split_points(&split_points, &sort_options)?; + Ok(Self { + on_columns, + split_points, + sort_options, + }) + } + + /// Get the columns used to compute Range partition IDs. + pub fn on_columns(&self) -> &[PhysicalExprRef] { + &self.on_columns + } + + /// Returns the Range split points used for routing. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Returns the per-key sort options used for routing. + pub fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } +} + +impl Display for RangeExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "range_partition") + } +} + +impl PhysicalExpr for RangeExpr { + fn children(&self) -> Vec<&PhysicalExprRef> { + self.on_columns.iter().collect() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<PhysicalExprRef>, + ) -> Result<PhysicalExprRef> { + assert_or_internal_err!( + children.len() == self.on_columns.len(), + "RangeExpr expected {} children, got {}", + self.on_columns.len(), + children.len() + ); + Ok(Arc::new(Self::try_new_parts( + children, + self.split_points.clone(), + self.sort_options.clone(), + )?)) + } + + fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { + Ok(DataType::UInt64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result<bool> { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + let arrays = evaluate_expressions_to_arrays(self.on_columns.iter(), batch)?; + let mut row_key_buffer = Vec::with_capacity(arrays.len()); + let mut partition_ids = Vec::with_capacity(batch.num_rows()); + for row_idx in 0..batch.num_rows() { + extract_row_at_idx_to_buf(&arrays, row_idx, &mut row_key_buffer)?; + partition_ids.push(range_partition_id( + &row_key_buffer, + &self.split_points, + &self.sort_options, + )? as u64); + } + Ok(ColumnarValue::Array(Arc::new(UInt64Array::from( + partition_ids, + )))) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "range_partition") + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result<Option<protobuf::PhysicalExprNode>> { + // Encode the raw ordered children: rebuilding a `LexOrdering` would + // deduplicate equivalent children after dynamic-filter remapping. + let sort_exprs = self + .on_columns + .iter() + .zip(&self.sort_options) + .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr), *options)) + .collect::<Vec<_>>(); + let sort_expr = sort_exprs_try_to_proto(&sort_exprs, ctx)?; + let split_point = self + .split_points + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::<Result<Vec<_>>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::RangeExpr( + protobuf::PhysicalRangeExprNode { + sort_expr, + split_point, + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl RangeExpr { + /// Reconstructs a [`RangeExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result<PhysicalExprRef> { + // Decode the raw ordered children for the same reason as `try_to_proto`. + let range_expr = match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::RangeExpr(expr)) => expr, + _ => return internal_err!("PhysicalExprNode is not a RangeExpr"), + }; + let sort_exprs = sort_exprs_try_from_proto(&range_expr.sort_expr, ctx)?; + let (on_columns, sort_options) = sort_exprs + .into_iter() + .map(|sort_expr| (sort_expr.expr, sort_expr.options)) + .unzip(); + let split_points = range_expr + .split_point + .iter() + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::<Result<Vec<_>>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Arc::new(Self::try_new_parts( + on_columns, + split_points, + sort_options, + )?)) + } +} + +fn range_partition_id( + row_key: &[ScalarValue], + split_points: &[SplitPoint], + sort_options: &[SortOptions], +) -> Result<usize> { + let mut low = 0; + let mut high = split_points.len(); + while low < high { + let mid = low + (high - low) / 2; + match compare_rows(row_key, split_points[mid].values(), sort_options)? { + Ordering::Less => high = mid, + Ordering::Equal | Ordering::Greater => low = mid + 1, + } + } + Ok(low) +} + +/// Computes `value % divisor` without division in the hot loop when `divisor` +/// is fixed for many values. +/// +/// Hash repartitioning computes a remainder for every row. Integer division is +/// relatively expensive, so this precomputes the strength-reduced form of the +/// divisor: powers of two use a bit mask, and other divisors use a reciprocal +/// multiply to recover the quotient and therefore the remainder. This is the +/// same invariant-divisor optimization compilers use for `%` by a constant. +#[derive(Debug, Clone, Copy)] +enum StrengthReducedU64 { + PowerOfTwo { mask: u64 }, + Reciprocal { divisor: u64, reciprocal: u128 }, +} + +impl StrengthReducedU64 { + fn new(divisor: u64) -> Self { + debug_assert!(divisor > 0); + + if divisor.is_power_of_two() { + Self::PowerOfTwo { mask: divisor - 1 } + } else { + Self::Reciprocal { + divisor, + // ceil(2^128 / divisor), computed without representing 2^128 + reciprocal: u128::MAX / u128::from(divisor) + 1, + } + } + } + + fn partition_indices(self, hash_buffer: &[u64], indices: &mut [Vec<u32>]) { + match self { + Self::PowerOfTwo { mask } => { + for (index, hash) in hash_buffer.iter().enumerate() { + indices[(*hash & mask) as usize].push(index as u32); + } + } + Self::Reciprocal { + divisor, + reciprocal, + } => { + for (index, hash) in hash_buffer.iter().enumerate() { + let quotient = Self::quotient(*hash, reciprocal); + let partition = *hash - quotient * divisor; + indices[partition as usize].push(index as u32); + } + } + } + } + + #[cfg(test)] + fn remainder(self, value: u64) -> u64 { + match self { + Self::PowerOfTwo { mask } => value & mask, + Self::Reciprocal { + divisor, + reciprocal, + } => value - Self::quotient(value, reciprocal) * divisor, + } + } + + #[inline] + fn quotient(value: u64, reciprocal: u128) -> u64 { + let reciprocal_low = reciprocal as u64; + let reciprocal_high = (reciprocal >> 64) as u64; + let low_product = u128::from(value) * u128::from(reciprocal_low); + let high_product = u128::from(value) * u128::from(reciprocal_high); + let carry = ((high_product & u128::from(u64::MAX)) + (low_product >> 64)) >> 64; + + ((high_product >> 64) + carry) as u64 + } +} + +impl BatchPartitioner { + /// Create a new [`BatchPartitioner`] for hash-based repartitioning. + /// + /// # Parameters + /// - `exprs`: Expressions used to compute the hash for each input row. + /// - `num_partitions`: Total number of output partitions. + /// - `timer`: Metric used to record time spent during repartitioning. + /// + /// The partition count is fixed for the lifetime of the partitioner, so this + /// precomputes a strength-reduced reducer for `hash % num_partitions`. + /// + /// # Errors + /// Returns an error if `num_partitions` is zero. + pub fn new_hash_partitioner( + exprs: Vec<Arc<dyn PhysicalExpr>>, + num_partitions: usize, + timer: metrics::Time, + ) -> Result<Self> { + if num_partitions == 0 { + return internal_err!("Hash repartition requires at least one partition"); + } + + Ok(Self { + state: BatchPartitionerState::Hash { + exprs, + partition_reducer: StrengthReducedU64::new(num_partitions as u64), + hash_buffer: vec![], + indices: vec![vec![]; num_partitions], + }, + timer, + }) + } + + /// Create a new [`BatchPartitioner`] for round-robin repartitioning. + /// + /// # Parameters + /// - `num_partitions`: Total number of output partitions. + /// - `timer`: Metric used to record time spent during repartitioning. + /// - `input_partition`: Index of the current input partition. + /// - `num_input_partitions`: Total number of input partitions. + /// + /// # Notes + /// The starting output partition is derived from the input partition + /// to avoid skew when multiple input partitions are used. + pub fn new_round_robin_partitioner( + num_partitions: usize, + timer: metrics::Time, + input_partition: usize, + num_input_partitions: usize, + ) -> Self { + Self { + state: BatchPartitionerState::RoundRobin { + num_partitions, + next_idx: (input_partition * num_partitions) / num_input_partitions, + }, + timer, + } + } + + /// Create a new [`BatchPartitioner`] for range-based repartitioning. + /// + /// # Parameters + /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions + /// - `timer`: Metric used to record time spent during repartitioning. + pub fn new_range_partitioner( + range_partitioning: &RangePartitioning, + timer: metrics::Time, + ) -> Self { + let ordering = range_partitioning.ordering().clone(); + let split_points = range_partitioning.split_points().to_vec(); + let num_partitions = range_partitioning.partition_count(); + let sort_options: Vec<SortOptions> = ordering.iter().map(|e| e.options).collect(); + + Self { + state: BatchPartitionerState::Range { + partition_buffer: Vec::with_capacity(ordering.len()), + ordering, + sort_options, + split_points, + indices: vec![vec![]; num_partitions], + }, + timer, + } + } + + /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. + /// + /// This is a convenience constructor that delegates to the specialized + /// hash, round-robin, or range constructors depending on the partitioning variant. + /// + /// # Parameters + /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range). + /// - `timer`: Metric used to record time spent during repartitioning. + /// - `input_partition`: Index of the current input partition. + /// - `num_input_partitions`: Total number of input partitions. + /// + /// # Errors + /// Returns an error if the provided partitioning scheme is not supported, + /// or if hash partitioning is requested with zero output partitions. + pub fn try_new( + partitioning: Partitioning, + timer: metrics::Time, + input_partition: usize, + num_input_partitions: usize, + ) -> Result<Self> { + match partitioning { + Partitioning::Hash(exprs, num_partitions) => { + Self::new_hash_partitioner(exprs, num_partitions, timer) + } + Partitioning::RoundRobinBatch(num_partitions) => { + Ok(Self::new_round_robin_partitioner( + num_partitions, + timer, + input_partition, + num_input_partitions, + )) + } + Partitioning::Range(range_repartitioning) => { + Ok(Self::new_range_partitioner(&range_repartitioning, timer)) + } + other => { + not_impl_err!("Unsupported repartitioning scheme {other:?}") + } + } + } + + /// Partition the provided [`RecordBatch`] into one or more partitioned [`RecordBatch`] + /// based on the [`Partitioning`] specified on construction + /// + /// `f` will be called for each partitioned [`RecordBatch`] with the corresponding + /// partition index. Any error returned by `f` will be immediately returned by this + /// function without attempting to publish further [`RecordBatch`] + /// + /// The time spent repartitioning, not including time spent in `f` will be recorded + /// to the [`metrics::Time`] provided on construction + pub fn partition<F>(&mut self, batch: RecordBatch, mut f: F) -> Result<()> + where + F: FnMut(usize, RecordBatch) -> Result<()>, + { + self.partition_iter(batch)?.try_for_each(|res| match res { + Ok((partition, batch)) => f(partition, batch), + Err(e) => Err(e), + }) + } + + /// Returns an iterator of `(partition_index, RecordBatch)` pairs for the given batch. + /// + /// This is useful for async consumers that want to separate CPU-bound partitioning + /// from I/O. For example, you can iterate results on the async side and send them + /// through a channel, while performing file I/O on a blocking task: + /// + /// ```ignore + /// for result in partitioner.partition_iter(batch)? { + /// let (partition, batch) = result?; + /// tx.send((partition, batch)).await?; + /// } + /// ``` + /// + /// The sync [`partition`](Self::partition) method is implemented on top of this. + pub fn partition_iter( + &mut self, + batch: RecordBatch, + ) -> Result<impl Iterator<Item = Result<(usize, RecordBatch)>> + Send + '_> { + let it: Box<dyn Iterator<Item = Result<(usize, RecordBatch)>> + Send> = + match &mut self.state { + BatchPartitionerState::RoundRobin { + num_partitions, + next_idx, + } => { + let idx = *next_idx; + *next_idx = (*next_idx + 1) % *num_partitions; + Box::new(std::iter::once(Ok((idx, batch)))) + } + BatchPartitionerState::Hash { + exprs, + partition_reducer, + hash_buffer, + indices, + } => { + // Tracking time required for distributing indexes across output partitions + let timer = self.timer.timer(); + + let arrays = + evaluate_expressions_to_arrays(exprs.as_slice(), &batch)?; + + hash_buffer.clear(); + hash_buffer.resize(batch.num_rows(), 0); + + create_hashes( + &arrays, + REPARTITION_RANDOM_STATE.random_state(), + hash_buffer, + )?; + + indices.iter_mut().for_each(|v| v.clear()); + + partition_reducer.partition_indices(hash_buffer, indices); + + // Finished building index-arrays for output partitions + timer.done(); + + let partitioned_batches = + Self::partition_grouped_take(&batch, indices, &self.timer)?; + + Box::new(partitioned_batches.into_iter()) + } + BatchPartitionerState::Range { + ordering, + sort_options, + split_points, + indices, + partition_buffer, + } => { + // Tracking time required for distributing indexes across output partitions + let timer = self.timer.timer(); + if split_points.is_empty() { + timer.done(); + Box::new(std::iter::once(Ok((0, batch)))) + } else { + let arrays = evaluate_expressions_to_arrays( + ordering.iter().map(|e| &e.expr), + &batch, + )?; + + indices.iter_mut().for_each(|v| v.clear()); + + Self::partition_range_indices( + &arrays, + split_points, + sort_options, + partition_buffer, + indices, + )?; + + // Finished building index-arrays for output partitions + timer.done(); + + let partitioned_batches = + Self::partition_grouped_take(&batch, indices, &self.timer)?; + + Box::new(partitioned_batches.into_iter()) + } + } + }; + + Ok(it) + } + + /// Groups input row indices by range partition. This populates `indices[p]` with the + /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`. + fn partition_range_indices( + arrays: &[Arc<dyn Array>], + split_points: &[SplitPoint], + sort_options: &[SortOptions], + row_key_buffer: &mut Vec<ScalarValue>, + indices: &mut [Vec<u32>], + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row + extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?; + + let partition = + range_partition_id(row_key_buffer, split_points, sort_options)?; + indices[partition].push(row_idx as u32) + } + + Ok(()) + } + + // return the number of output partitions + fn num_partitions(&self) -> usize { + match &self.state { + BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions, + BatchPartitionerState::Hash { indices, .. } + | BatchPartitionerState::Range { indices, .. } => indices.len(), + } + } + + /// Build repartitioned hash/range output batches using one `take` per input batch. + /// + /// The routers first fills one index vector per output partition. This method + /// concatenates those index vectors, performs one grouped `take_arrays`, and + /// then returns each output partition as a slice of the reordered batch. + /// + /// For example, given partition indices: + /// + /// ```text + /// partition 0: [2, 5] + /// partition 1: [] + /// partition 2: [0, 3, 4] + /// ``` + /// + /// this method takes rows in `[2, 5, 0, 3, 4]` order once, then returns + /// `partition 0 = slice(0, 2)` and `partition 2 = slice(2, 3)`. + fn partition_grouped_take( + batch: &RecordBatch, + indices: &mut [Vec<u32>], + timer: &metrics::Time, + ) -> Result<Vec<Result<(usize, RecordBatch)>>> { + let mut partition_ranges = Vec::with_capacity(indices.len()); + let mut reordered_indices = Vec::with_capacity(batch.num_rows()); + + for (partition, p_indices) in indices.iter_mut().enumerate() { + if p_indices.is_empty() { + continue; + } + + let start = reordered_indices.len(); + reordered_indices.extend_from_slice(p_indices); + partition_ranges.push((partition, start, p_indices.len())); + p_indices.clear(); + } + + if reordered_indices.is_empty() { + return Ok(vec![]); + } + + let batches = { + let _timer = timer.timer(); + let indices_array: PrimitiveArray<UInt32Type> = reordered_indices.into(); + let columns = take_arrays(batch.columns(), &indices_array, None)?; + + let mut options = RecordBatchOptions::new(); + options = options.with_row_count(Some(indices_array.len())); + let reordered_batch = + RecordBatch::try_new_with_options(batch.schema(), columns, &options)?; + + partition_ranges + .into_iter() + .map(|(partition, start, len)| { + Ok((partition, reordered_batch.slice(start, len))) + }) + .collect() + }; + + Ok(batches) + } +} + +/// Maps `N` input partitions to `M` output partitions based on a +/// [`Partitioning`] scheme. +/// +/// # Background +/// +/// DataFusion, like most other commercial systems, with the +/// notable exception of DuckDB, uses the "Exchange Operator" based +/// approach to parallelism which works well in practice given +/// sufficient care in implementation. +/// +/// DataFusion's planner picks the target number of partitions and +/// then [`RepartitionExec`] redistributes [`RecordBatch`]es to that number +/// of output partitions. +/// +/// For example, given `target_partitions=3` (trying to use 3 cores) +/// but scanning an input with 2 partitions, `RepartitionExec` can be +/// used to get 3 even streams of `RecordBatch`es +/// +/// +/// ```text +/// ▲ ▲ ▲ +/// │ │ │ +/// │ │ │ +/// │ │ │ +/// ┌───────────────┐ ┌───────────────┐ ┌───────────────┐ +/// │ GroupBy │ │ GroupBy │ │ GroupBy │ +/// │ (Partial) │ │ (Partial) │ │ (Partial) │ +/// └───────────────┘ └───────────────┘ └───────────────┘ +/// ▲ ▲ ▲ +/// └──────────────────┼──────────────────┘ +/// │ +/// ┌─────────────────────────┐ +/// │ RepartitionExec │ +/// │ (hash/round robin) │ +/// └─────────────────────────┘ +/// ▲ ▲ +/// ┌───────────┘ └───────────┐ +/// │ │ +/// │ │ +/// .─────────. .─────────. +/// ,─' '─. ,─' '─. +/// ; Input : ; Input : +/// : Partition 0 ; : Partition 1 ; +/// ╲ ╱ ╲ ╱ +/// '─. ,─' '─. ,─' +/// `───────' `───────' +/// ``` +/// +/// # Error Handling +/// +/// If any of the input partitions return an error, the error is propagated to +/// all output partitions and inputs are not polled again. +/// +/// # Output Ordering +/// +/// If more than one stream is being repartitioned, the output will be some +/// arbitrary interleaving (and thus unordered) unless +/// [`Self::with_preserve_order`] specifies otherwise. +/// +/// # Batch coalescing +/// +/// Repartitioning one [`RecordBatch`] implies creating multiple smaller batches, potentially +/// as many as the number of output partitions. [`RepartitionExec`] makes sure that the returned +/// batches adhere to the configured `datafusion.execution.batch_size` for efficient operations, +/// and for that, it will automatically coalesce batches right after repartitioning for bounded +/// inputs. Coalescing is skipped for unbounded inputs so partial batches are emitted promptly. +/// +/// For this, one shared [`LimitedBatchCoalescer`] per output partition is used: +/// +/// ```text +/// ┌───┐ ┌───┐ +/// ┌─▶│ │────────▶.───────────. │ │ ┌──────────────────┐ +/// │ └───┘ ┌───┐ ( Coalescer 0 )──▶ ├───┤ ───▶│ Output 0 │ +/// │┌──────▶│ │──▶`───────────' │ │ └──────────────────┘ +/// ││ └───┘ └───┘ +/// ┌──────────────────┐ ││ ┌──────────────────┐ +/// │BatchPartitioner 0│─┘│ │ Output 1 │ +/// └──────────────────┘ │ └──────────────────┘ +/// │ +/// ┌──────────────────┐ │ ... ┌──────────────────┐ +/// │BatchPartitioner 1│──┘ │ Output 2 │ +/// └──────────────────┘ └──────────────────┘ +/// +/// ┌──────────────────┐ +/// │ Output 3 │ +/// └──────────────────┘ +/// ``` +/// +/// # Spilling Architecture +/// +/// RepartitionExec uses [`SpillPool`](crate::spill::spill_pool) channels to handle +/// memory pressure during repartitioning. Each (input partition, output partition) +/// pair gets its own SpillPool channel for FIFO ordering. +/// +/// ```text +/// Input Partitions (N) Output Partitions (M) +/// ──────────────────── ───────────────────── +/// +/// Input 0 ──┐ ┌──▶ Output 0 +/// │ ┌──────────────┐ │ +/// ├─▶│ SpillPool │────┤ +/// │ │ [In0→Out0] │ │ +/// Input 1 ──┤ └──────────────┘ ├──▶ Output 1 +/// │ │ +/// │ ┌──────────────┐ │ +/// ├─▶│ SpillPool │────┤ +/// │ │ [In1→Out0] │ │ +/// Input 2 ──┤ └──────────────┘ ├──▶ Output 2 +/// │ │ +/// │ ... (N×M SpillPools total) +/// │ │ +/// │ ┌──────────────┐ │ +/// └─▶│ SpillPool │────┘ +/// │ [InN→OutM] │ +/// └──────────────┘ +/// +/// Each SpillPool maintains FIFO order for its (input, output) pair. +/// See `RepartitionBatch` for details on the memory/spill decision logic. +/// ``` +/// +/// # Footnote +/// +/// The "Exchange Operator" was first described in the 1989 paper +/// [Encapsulation of parallelism in the Volcano query processing +/// system Paper](https://dl.acm.org/doi/pdf/10.1145/93605.98720) +/// which uses the term "Exchange" for the concept of repartitioning +/// data across threads. +/// +/// For more background, please also see the [Optimizing Repartitions in DataFusion] blog. +/// +/// [Optimizing Repartitions in DataFusion]: https://datafusion.apache.org/blog/2025/12/15/avoid-consecutive-repartitions +#[derive(Debug, Clone)] +pub struct RepartitionExec { + /// Input execution plan + input: Arc<dyn ExecutionPlan>, + /// Inner state that is initialized when the parent calls .execute() on this node + /// and consumed as soon as the parent starts consuming this node. + state: Arc<Mutex<RepartitionExecState>>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Boolean flag to decide whether to preserve ordering. If true means + /// `SortPreservingRepartitionExec`, false means `RepartitionExec`. + preserve_order: bool, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +#[derive(Debug, Clone)] +struct RepartitionMetrics { + /// Time in nanos to execute child operator and fetch batches + fetch_time: metrics::Time, + /// Repartitioning elapsed time in nanos + repartition_time: metrics::Time, + /// Time in nanos for sending resulting batches to channels. + /// + /// One metric per output partition. + send_time: Vec<metrics::Time>, +} + +impl RepartitionMetrics { + pub fn new( + input_partition: usize, + num_output_partitions: usize, + metrics: &ExecutionPlanMetricsSet, + ) -> Self { + // Time in nanos to execute child operator and fetch batches + let fetch_time = + MetricBuilder::new(metrics).subset_time("fetch_time", input_partition); + + // Time in nanos to perform repartitioning + let repartition_time = + MetricBuilder::new(metrics).subset_time("repartition_time", input_partition); + + // Time in nanos for sending resulting batches to channels + let send_time = (0..num_output_partitions) + .map(|output_partition| { + let label = + metrics::Label::new("outputPartition", output_partition.to_string()); + MetricBuilder::new(metrics) + .with_label(label) + .subset_time("send_time", input_partition) + }) + .collect(); + + Self { + fetch_time, + repartition_time, + send_time, + } + } +} + +impl RepartitionExec { + /// Input execution plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Partitioning scheme to use + pub fn partitioning(&self) -> &Partitioning { + &self.cache.partitioning + } + + /// Get preserve_order flag of the RepartitionExec + /// `true` means `SortPreservingRepartitionExec`, `false` means `RepartitionExec` + pub fn preserve_order(&self) -> bool { + self.preserve_order + } + + /// Get name used to display this Exec + pub fn name(&self) -> &str { + "RepartitionExec" + } +} + +impl DisplayAs for RepartitionExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + let input_partition_count = self.input.output_partitioning().partition_count(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "{}: partitioning={}, input_partitions={}", + self.name(), + self.partitioning(), + input_partition_count, + )?; + + if self.preserve_order { + write!(f, ", preserve_order=true")?; + } else if input_partition_count <= 1 + && self.input.output_ordering().is_some() + { + // Make it explicit that repartition maintains sortedness for a single input partition even + // when `preserve_sort order` is false + write!(f, ", maintains_sort_order=true")?; + } + + if let Some(sort_exprs) = self.sort_exprs() { + write!(f, ", sort_exprs={}", sort_exprs.clone())?; + } + Ok(()) + } + DisplayFormatType::TreeRender => { + writeln!(f, "partitioning_scheme={}", self.partitioning(),)?; + let output_partition_count = self.partitioning().partition_count(); + let input_to_output_partition_str = + format!("{input_partition_count} -> {output_partition_count}"); + writeln!( + f, + "partition_count(in->out)={input_to_output_partition_str}" + )?; + + if self.preserve_order { + writeln!(f, "preserve_order={}", self.preserve_order)?; + } + Ok(()) + } + } + } +} + +impl ExecutionPlan for RepartitionExec { + fn name(&self) -> &'static str { + "RepartitionExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + match self.partitioning() { + Partitioning::Hash(exprs, _) => crate::apply_expression_roots(exprs, f), + Partitioning::Range(range) => crate::apply_expression_roots( + range.ordering().iter().map(|sort_expr| &sort_expr.expr), + f, + ), + _ => Ok(TreeNodeRecursion::Continue), + } + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut repartition = RepartitionExec::try_new( + children.swap_remove(0), + self.partitioning().clone(), + )?; + if self.preserve_order { + repartition = repartition.with_preserve_order(); + } + Ok(Arc::new(repartition)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![matches!(self.partitioning(), Partitioning::Hash(_, _))] + } + + fn maintains_input_order(&self) -> Vec<bool> { + Self::maintains_input_order_helper(self.input(), self.preserve_order) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start {}::execute for partition: {}", + self.name(), + partition + ); + + let spill_metrics = SpillMetrics::new(&self.metrics, partition); + + let input = Arc::clone(&self.input); + let partitioning = self.partitioning().clone(); + let metrics = self.metrics.clone(); + let preserve_order = self.sort_exprs().is_some(); + let name = self.name().to_owned(); + let schema = self.schema(); + let schema_captured = Arc::clone(&schema); + + let spill_manager = SpillManager::new( + Arc::clone(&context.runtime_env()), + spill_metrics, + input.schema(), + ); + + // Get existing ordering to use for merging + let sort_exprs = self.sort_exprs().cloned(); + + let state = Arc::clone(&self.state); + if let Some(mut state) = state.try_lock() { + state.ensure_input_streams_initialized( + &input, + &metrics, + partitioning.partition_count(), + &context, + )?; + } + + let num_input_partitions = input.output_partitioning().partition_count(); + + let stream = futures::stream::once(async move { + // lock scope + let (rx, reservation, spill_readers, abort_helper) = { + // lock mutexes + let mut state = state.lock(); + let state = state.consume_input_streams( + &input, + &metrics, + &partitioning, + preserve_order, + &name, + &context, + spill_manager.clone(), + )?; + + // now return stream for the specified *output* partition which will + // read from the channel + let PartitionChannels { + rx, + reservation, + spill_readers, + .. + } = state + .channels + .remove(&partition) + .expect("partition not used yet"); + + ( + rx, + reservation, + spill_readers, + Arc::clone(&state.abort_helper), + ) + }; + + trace!( + "Before returning stream in {name}::execute for partition: {partition}" + ); + + if preserve_order { + // Store streams from all the input partitions: + // Each input partition gets its own spill reader to maintain proper FIFO ordering + // + // Pass None for metrics here — these intermediate streams feed into + // StreamingMerge which is the actual output. Only the merge's + // BaselineMetrics should contribute to the operator's reported + // output_rows. Without this, every row would be counted twice + // (once by PerPartitionStream, once by StreamingMerge). + let input_streams = rx + .into_iter() + .zip(spill_readers) + .map(|(receiver, spill_stream)| { + // In preserve_order mode, each receiver corresponds to exactly one input partition + Box::pin(PerPartitionStream::new( + Arc::clone(&schema_captured), + receiver, + Arc::clone(&abort_helper), + Arc::clone(&reservation), + spill_stream, + 1, // Each receiver handles one input partition + None, + )) as SendableRecordBatchStream + }) + .collect::<Vec<_>>(); + // Note that receiver size (`rx.len()`) and `num_input_partitions` are same. + + // Merge streams (while preserving ordering) coming from + // input partitions to this partition: + let fetch = None; + let merge_reservation = + MemoryConsumer::new(format!("{name}[Merge {partition}]")) + .register(context.memory_pool()); + StreamingMergeBuilder::new() + .with_streams(input_streams) + .with_schema(schema_captured) + .with_expressions(&sort_exprs.unwrap()) + .with_metrics(BaselineMetrics::new(&metrics, partition)) + .with_batch_size(context.session_config().batch_size()) + .with_fetch(fetch) + .with_reservation(merge_reservation) + .with_spill_manager(spill_manager) + .build() + } else { + // Non-preserve-order case: single input stream, so use the first spill reader + let spill_stream = spill_readers + .into_iter() + .next() + .expect("at least one spill reader should exist"); + + Ok(Box::pin(PerPartitionStream::new( + schema_captured, + rx.into_iter() + .next() + .expect("at least one receiver should exist"), + abort_helper, + reservation, + spill_stream, + num_input_partitions, + Some(BaselineMetrics::new(&metrics, partition)), + )) as SendableRecordBatchStream) + } + }) + .try_flatten(); + let stream = RecordBatchStreamAdapter::new(schema, stream); + Ok(Box::pin(stream)) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(None)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if args.partition().is_some() { + let partition_count = self.partitioning().partition_count(); + // `StatisticsContext::compute` validates the partition index against + // this same count before calling, so it is non-zero here; guard + // defensively against a direct call so the division below cannot + // divide by zero + assert_or_internal_err!( + partition_count > 0, + "RepartitionExec statistics requested for a partition but the partition count is 0" + ); + + let mut stats = input_stats[0].as_ref().clone(); + + // Distribute statistics across partitions + stats.num_rows = stats + .num_rows + .get_value() + .map(|rows| Precision::Inexact(rows / partition_count)) + .unwrap_or(Precision::Absent); + stats.total_byte_size = stats + .total_byte_size + .get_value() + .map(|bytes| Precision::Inexact(bytes / partition_count)) + .unwrap_or(Precision::Absent); + + // Make all column stats unknown + stats.column_statistics = stats + .column_statistics + .iter() + .map(|_| ColumnStatistics::new_unknown()) + .collect(); + + Ok(Arc::new(stats)) + } else { + Ok(Arc::clone(&input_stats[0])) + } + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection does not narrow the schema, we should not try to push it down. + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + + // If pushdown is not beneficial or applicable, break it. + if projection.benefits_from_input_partitioning()[0] + || !all_columns(projection.expr()) + { + return Ok(None); + } + + let new_projection = make_with_child(projection, self.input())?; + + let new_partitioning = match self.partitioning() { + Partitioning::Hash(partitions, size) => { + let mut new_partitions = vec![]; + for partition in partitions { + let Some(new_partition) = + update_expr(partition, projection.expr(), false)? + else { + return Ok(None); + }; + new_partitions.push(new_partition); + } + Partitioning::Hash(new_partitions, *size) + } + Partitioning::Range(range_partitioning) => { + // Rewrite range key expressions through the projection. + let mut sort_exprs = + Vec::with_capacity(range_partitioning.ordering().len()); + for sort_expr in range_partitioning.ordering() { + let Some(new_expr) = + update_expr(&sort_expr.expr, projection.expr(), false)? + else { + return Ok(None); + }; + sort_exprs.push(PhysicalSortExpr::new(new_expr, sort_expr.options)); + } + + let Some(ordering) = LexOrdering::new(sort_exprs) else { + return internal_err!( + "failed to create LexOrdering for range partitioning" + ); + }; + + Partitioning::Range(RangePartitioning::try_new( + ordering, + range_partitioning.split_points().to_vec(), + )?) + } + others => others.clone(), + }; + + Ok(Some(Arc::new(RepartitionExec::try_new( + new_projection, + new_partitioning, + )?))) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + FilterDescription::from_children(parent_filters, &self.children()) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + Ok(FilterPushdownPropagation::if_all(child_pushdown_result)) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + // RepartitionExec only maintains input order if preserve_order is set + // or if there's only one partition + if !self.maintains_input_order()[0] { + return Ok(SortOrderPushdownResult::Unsupported); + } + + // Delegate to the child and wrap with a new RepartitionExec + self.input.try_pushdown_sort(order)?.try_map(|new_input| { + let mut new_repartition = + RepartitionExec::try_new(new_input, self.partitioning().clone())?; + if self.preserve_order { + new_repartition = new_repartition.with_preserve_order(); + } + Ok(Arc::new(new_repartition) as Arc<dyn ExecutionPlan>) + }) + } + + fn repartitioned( + &self, + target_partitions: usize, + _config: &ConfigOptions, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + use Partitioning::*; + let mut new_properties = PlanProperties::clone(&self.cache); + new_properties.partitioning = match new_properties.partitioning { + RoundRobinBatch(_) => RoundRobinBatch(target_partitions), + Hash(hash, _) => Hash(hash, target_partitions), + Range(_) => { + // Number of partitions is constrained by the split points and cannot be changed + return Ok(None); + } + UnknownPartitioning(_) => UnknownPartitioning(target_partitions), + }; + Ok(Some(Arc::new(Self { + input: Arc::clone(&self.input), + state: Arc::clone(&self.state), + metrics: self.metrics.clone(), + preserve_order: self.preserve_order, + cache: new_properties.into(), + }))) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<protobuf::PhysicalPlanNode>> { + let input = ctx.encode_child(self.input())?; + + let partitioning = self.partitioning().try_to_proto(&ctx.expr_ctx())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Repartition(Box::new( + protobuf::RepartitionExecNode { + input: Some(Box::new(input)), + partitioning: Some(partitioning), + preserve_order: self.preserve_order(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl RepartitionExec { + /// Reconstruct a [`RepartitionExec`] from its protobuf representation. + pub fn try_from_proto( + node: &protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let repart = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Repartition, + "RepartitionExec", + ); + let input = ctx.decode_required_child( + repart.input.as_deref(), + "RepartitionExec", + "input", + )?; + let input_schema = input.schema(); + + let partitioning = repart + .partitioning + .as_ref() + .map(|partitioning| { + Partitioning::try_from_proto( + partitioning, + &ctx.expr_ctx(input_schema.as_ref()), + ) + }) + .transpose()? + .flatten() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "RepartitionExec is missing required field 'partitioning'" + ) + })?; + + let mut repart_exec = RepartitionExec::try_new(input, partitioning)?; + if repart.preserve_order { + repart_exec = repart_exec.with_preserve_order(); + } + Ok(Arc::new(repart_exec)) + } +} + +impl RepartitionExec { + /// Create a new RepartitionExec, that produces output `partitioning`, and + /// does not preserve the order of the input (see [`Self::with_preserve_order`] + /// for more details) + pub fn try_new( + input: Arc<dyn ExecutionPlan>, + partitioning: Partitioning, + ) -> Result<Self> { + let preserve_order = false; + let cache = Self::compute_properties(&input, partitioning, preserve_order); + Ok(RepartitionExec { + input, + state: Default::default(), + metrics: ExecutionPlanMetricsSet::new(), + preserve_order, + cache: Arc::new(cache), + }) + } + + fn maintains_input_order_helper( + input: &Arc<dyn ExecutionPlan>, + preserve_order: bool, + ) -> Vec<bool> { + // We preserve ordering when repartition is order preserving variant or input partitioning is 1 + vec![preserve_order || input.output_partitioning().partition_count() <= 1] + } + + fn eq_properties_helper( + input: &Arc<dyn ExecutionPlan>, + preserve_order: bool, + ) -> EquivalenceProperties { + // Equivalence Properties + let mut eq_properties = input.equivalence_properties().clone(); + // If the ordering is lost, reset the ordering equivalence class: + if !Self::maintains_input_order_helper(input, preserve_order)[0] { + eq_properties.clear_orderings(); + } + // When there are more than one input partitions, they will be fused at the output. + // Therefore, remove per partition constants. + if input.output_partitioning().partition_count() > 1 { + eq_properties.clear_per_partition_constants(); + } + eq_properties + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + partitioning: Partitioning, + preserve_order: bool, + ) -> PlanProperties { + PlanProperties::new( + Self::eq_properties_helper(input, preserve_order), + partitioning, + input.pipeline_behavior(), + input.boundedness(), + ) + .with_scheduling_type(SchedulingType::Cooperative) + .with_evaluation_type(EvaluationType::Eager) + } + + /// Specify if this repartitioning operation should preserve the order of + /// rows from its input when producing output. Preserving order is more + /// expensive at runtime, so should only be set if the output of this + /// operator can take advantage of it. + /// + /// If the input is not ordered, or has only one partition, this is a no op, + /// and the node remains a `RepartitionExec`. + pub fn with_preserve_order(mut self) -> Self { + self.preserve_order = + // If the input isn't ordered, there is no ordering to preserve + self.input.output_ordering().is_some() && + // if there is only one input partition, merging is not required + // to maintain order + self.input.output_partitioning().partition_count() > 1; + let eq_properties = Self::eq_properties_helper(&self.input, self.preserve_order); + Arc::make_mut(&mut self.cache).set_eq_properties(eq_properties); + self + } + + /// Return the sort expressions that are used to merge + fn sort_exprs(&self) -> Option<&LexOrdering> { + if self.preserve_order { + self.input.output_ordering() + } else { + None + } + } + + /// Pulls data from the specified input plan, feeding it to the + /// output partitions based on the desired partitioning + /// + /// `output_channels` holds the output sending channels for each output partition + async fn pull_from_input( + mut stream: SendableRecordBatchStream, + mut output_channels: HashMap<usize, OutputChannel>, + partitioning: Partitioning, + metrics: RepartitionMetrics, + input_partition: usize, + num_input_partitions: usize, + ) -> Result<()> { + let mut partitioner = BatchPartitioner::try_new( + partitioning, + metrics.repartition_time.clone(), + input_partition, + num_input_partitions, + )?; + + // While there are still outputs to send to, keep pulling inputs + let mut batches_until_yield = partitioner.num_partitions(); + while !output_channels.is_empty() { + // fetch the next batch + let timer = metrics.fetch_time.timer(); + let result = stream.next().await; + timer.done(); + + // Input is done + let batch = match result { + Some(result) => result?, + None => break, + }; + + // Handle empty batch + if batch.num_rows() == 0 { + continue; + } + + for res in partitioner.partition_iter(batch)? { + let (partition, batch) = res?; + + let timer = metrics.send_time[partition].timer(); + // if there is still a receiver, send to it + if let Some(output_channel) = output_channels.get_mut(&partition) { + for batch in output_channel.coalesce(batch)? { + if output_channel.send(batch).await.is_err() { + // If the other end has hung up, it was an early shutdown (e.g. LIMIT) + // so ignore this channel from now on. + output_channels.remove(&partition); + break; + } + } + } + timer.done(); + } + + // If the input stream is endless, we may spin forever and + // never yield back to tokio. See + // https://github.com/apache/datafusion/issues/5278. + // + // However, yielding on every batch causes a bottleneck + // when running with multiple cores. See + // https://github.com/apache/datafusion/issues/6290 + // + // Thus, heuristically yield after producing num_partition + // batches + // + // In round robin this is ideal as each input will get a + // new batch. In hash partitioning it may yield too often + // on uneven distributions even if some partition can not + // make progress, but parallelism is going to be limited + // in that case anyways + if batches_until_yield == 0 { + tokio::task::yield_now().await; + batches_until_yield = partitioner.num_partitions(); + } else { + batches_until_yield -= 1; + } + } + + // End of input for this task. For each output partition we still + // have a channel to, decrement the active-senders counter; whoever + // sees the count drop to zero is the last input task and must + // finalize the shared coalescer and ship its residual. + for (_, output_channel) in output_channels.drain() { + output_channel.finalize().await?; + } + + // Spill writers will auto-finalize when dropped + // No need for explicit flush + Ok(()) + } + + /// Waits for `input_task` which is consuming one of the inputs to + /// complete. Upon each successful completion, sends a `None` to + /// each of the output tx channels to signal one of the inputs is + /// complete. Upon error, propagates the errors to all output tx + /// channels. + async fn wait_for_task( + input_task: SpawnedTask<Result<()>>, + txs: HashMap<usize, DistributionSender<MaybeBatch>>, + ) { + // wait for completion, and propagate error + // note we ignore errors on send (.ok) as that means the receiver has already shutdown. + + match input_task.join().await { + // Error in joining task + Err(e) => { + let e = Arc::new(e); + + for (_, tx) in txs { + let err = Err(DataFusionError::Context( + "Join Error".to_string(), + Box::new(DataFusionError::External(Box::new(Arc::clone(&e)))), + )); + tx.send(Some(err)).await.ok(); + } + } + // Error from running input task + Ok(Err(e)) => { + // send the same Arc'd error to all output partitions + let e = Arc::new(e); + + for (_, tx) in txs { + // wrap it because need to send error to all output partitions + let err = Err(DataFusionError::from(&e)); + tx.send(Some(err)).await.ok(); + } + } + // Input task completed successfully + Ok(Ok(())) => { + // notify each output partition that this input partition has no more data + for (_partition, tx) in txs { + tx.send(None).await.ok(); + } + } + } + } +} + +/// State for tracking whether we're reading from memory channel or spill stream. +/// +/// This state machine ensures proper ordering when batches are mixed between memory +/// and spilled storage. When a [`RepartitionBatch::Spilled`] marker is received, +/// the stream must block on the spill stream until the corresponding batch arrives. +/// +/// # State Machine +/// +/// ```text +/// ┌─────────────────┐ +/// ┌───▶│ ReadingMemory │◀───┐ +/// │ └────────┬────────┘ │ +/// │ │ │ +/// │ Poll channel │ +/// │ │ │ +/// │ ┌──────────┼─────────────┐ +/// │ │ │ │ +/// │ ▼ ▼ │ +/// │ Memory Spilled │ +/// Got batch │ batch marker │ +/// from spill │ │ │ │ +/// │ │ ▼ │ +/// │ │ ┌──────────────────┐ │ +/// │ │ │ ReadingSpilled │ │ +/// │ │ └────────┬─────────┘ │ +/// │ │ │ │ +/// │ │ Poll spill_stream │ +/// │ │ │ │ +/// │ │ ▼ │ +/// │ │ Get batch │ +/// │ │ │ │ +/// └──┴───────────┴────────────┘ +/// │ +/// ▼ +/// Return batch +/// (Order preserved within +/// (input, output) pair) +/// ``` +/// +/// The transition to `ReadingSpilled` blocks further channel polling to maintain +/// FIFO ordering - we cannot read the next item from the channel until the spill +/// stream provides the current batch. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum StreamState { + /// Reading from the memory channel (normal operation) + ReadingMemory, + /// Waiting for a spilled batch from the spill stream. + /// Must not poll channel until spilled batch is received to preserve ordering. + ReadingSpilled, +} + +/// This struct converts a receiver to a stream. +/// Receiver receives data on an SPSC channel. +struct PerPartitionStream { + /// Schema wrapped by Arc + schema: SchemaRef, + + /// channel containing the repartitioned batches + receiver: DistributionReceiver<MaybeBatch>, + + /// Handle to ensure background tasks are killed when no longer needed. + _drop_helper: Arc<Vec<SpawnedTask<()>>>, + + /// Memory reservation. + reservation: SharedMemoryReservation, + + /// Infinite stream for reading from the spill pool + spill_stream: SendableRecordBatchStream, + + /// Internal state indicating if we are reading from memory or spill stream + state: StreamState, + + /// Number of input partitions that have not yet finished. + /// In non-preserve-order mode, multiple input partitions send to the same channel, + /// each sending None when complete. We must wait for all of them. + remaining_partitions: usize, + + /// Execution metrics (None in preserve-order mode where StreamingMerge owns the metrics) + baseline_metrics: Option<BaselineMetrics>, +} + +impl PerPartitionStream { + fn new( + schema: SchemaRef, + receiver: DistributionReceiver<MaybeBatch>, + drop_helper: Arc<Vec<SpawnedTask<()>>>, + reservation: SharedMemoryReservation, + spill_stream: SendableRecordBatchStream, + num_input_partitions: usize, + baseline_metrics: Option<BaselineMetrics>, + ) -> Self { + Self { + schema, + receiver, + _drop_helper: drop_helper, + reservation, + spill_stream, + state: StreamState::ReadingMemory, + remaining_partitions: num_input_partitions, + baseline_metrics, + } + } + + fn poll_next_inner( + self: &mut Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + use futures::StreamExt; + let elapsed = self + .baseline_metrics + .as_ref() + .map(|m| m.elapsed_compute().clone()); + let _timer = elapsed.as_ref().map(|t| t.timer()); + + loop { + match self.state { + StreamState::ReadingMemory => { + // Poll the memory channel for next message + let value = match self.receiver.recv().poll_unpin(cx) { + Poll::Ready(v) => v, + Poll::Pending => { + // Nothing from channel, wait + return Poll::Pending; + } + }; + + match value { + Some(Some(v)) => match v { + Ok(RepartitionBatch::Memory(batch)) => { + // Release memory and return batch + self.reservation.shrink(batch.get_array_memory_size()); + return Poll::Ready(Some(Ok(batch))); + } + Ok(RepartitionBatch::Spilled) => { + // Batch was spilled, transition to reading from spill stream + // We must block on spill stream until we get the batch + // to preserve ordering + self.state = StreamState::ReadingSpilled; + continue; + } + Err(e) => { + return Poll::Ready(Some(Err(e))); + } + }, + Some(None) => { + // One input partition finished + self.remaining_partitions -= 1; + if self.remaining_partitions == 0 { + // All input partitions finished + return Poll::Ready(None); + } + // Continue to poll for more data from other partitions + continue; + } + None => { + // Channel closed unexpectedly + return Poll::Ready(None); + } + } + } + StreamState::ReadingSpilled => { + // Poll spill stream for the spilled batch + match self.spill_stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + self.state = StreamState::ReadingMemory; + return Poll::Ready(Some(Ok(batch))); + } + Poll::Ready(Some(Err(e))) => { + return Poll::Ready(Some(Err(e))); + } + Poll::Ready(None) => { + // Spill stream ended — release its resources before + // we go back to draining the memory channel. + let spill_schema = self.spill_stream.schema(); + self.spill_stream = + Box::pin(EmptyRecordBatchStream::new(spill_schema)); + self.state = StreamState::ReadingMemory; + } + Poll::Pending => { + // Spilled batch not ready yet, must wait + // This preserves ordering by blocking until spill data arrives + return Poll::Pending; + } + } + } + } + } + } +} + +impl Stream for PerPartitionStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let poll = self.poll_next_inner(cx); + if let Some(metrics) = &self.baseline_metrics { + metrics.record_poll(poll) + } else { + poll + } + } +} + +impl RecordBatchStream for PerPartitionStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExpr; + use crate::streaming::{PartitionStream, StreamingTableExec}; + use crate::test::TestMemoryExec; + use crate::{ + test::{ + assert_is_pending, + exec::{ + BarrierExec, BlockingExec, ErrorExec, MockExec, + assert_strong_count_converges_to_zero, + }, + }, + {collect, expressions::col}, + }; + + use arrow::array::{ArrayRef, StringArray, UInt32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::ScalarValue; + use datafusion_common::cast::{as_string_array, as_uint32_array}; + use datafusion_common::exec_err; + use datafusion_common::test_util::batches_to_sort_string; + use datafusion_common_runtime::JoinSet; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; + use insta::assert_snapshot; + + #[derive(Debug)] + struct UnboundedTestPartition { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for UnboundedTestPartition { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream { + let stream = futures::stream::iter([Ok(self.batch.clone())]) + .chain(futures::stream::pending()); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + )) + } + } + + #[test] + fn range_expr_preserves_duplicate_remapped_children() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let sort_options = [SortOptions::new(false, false), SortOptions::new(true, true)]; + let split_points = vec![SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(20)), + ])]; + let range_partitioning = RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, sort_options[0]), + PhysicalSortExpr::new(col("b", &schema)?, sort_options[1]), + ] + .into(), + split_points.clone(), + )?; + let expr = Arc::new(RangeExpr::try_new( + vec![col("a", &schema)?, col("b", &schema)?], + &range_partitioning, + )?); + let remapped = col("a", &schema)?; + let rewritten = + expr.with_new_children(vec![Arc::clone(&remapped), Arc::clone(&remapped)])?; + + let rewritten = rewritten + .downcast_ref::<RangeExpr>() + .expect("rewritten expression should remain a RangeExpr"); + assert_eq!(rewritten.on_columns().len(), 2); + assert!(Arc::ptr_eq( + &rewritten.on_columns()[0], + &rewritten.on_columns()[1] + )); + assert_eq!(rewritten.sort_options(), sort_options); + assert_eq!(rewritten.split_points(), split_points); + + Ok(()) + } + + #[test] + fn strength_reduced_u64_remainder_matches_modulo() { + let divisors = [ + 1, + 2, + 3, + 4, + 5, + 7, + 8, + 10, + 16, + 31, + 32, + 63, + 64, + 65, + 97, + u64::from(u32::MAX), + u64::from(u32::MAX) + 1, + 1_u64 << 32, + (1_u64 << 63) - 1, + 1_u64 << 63, + u64::MAX - 1, + u64::MAX, + ]; + let values = [ + 0, + 1, + 2, + 3, + 4, + 5, + 31, + 32, + 33, + 63, + 64, + 65, + u64::from(u32::MAX) - 1, + u64::from(u32::MAX), + u64::from(u32::MAX) + 1, + (1_u64 << 32) - 1, + 1_u64 << 32, + (1_u64 << 32) + 1, + (1_u64 << 63) - 1, + 1_u64 << 63, + (1_u64 << 63) + 1, + u64::MAX - 1, + u64::MAX, + ]; + + for divisor in divisors { + let reducer = StrengthReducedU64::new(divisor); + for value in values { + assert_eq!( + reducer.remainder(value), + value % divisor, + "value={value} divisor={divisor}" + ); + } + + let mut value = 0x1234_5678_9abc_def0 ^ divisor; + for _ in 0..10_000 { + value = value + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + assert_eq!( + reducer.remainder(value), + value % divisor, + "value={value} divisor={divisor}" + ); + } + } + } + + #[test] + fn hash_partitioner_requires_nonzero_partitions() { + let metrics = ExecutionPlanMetricsSet::new(); + let timer = MetricBuilder::new(&metrics).subset_time("test", 0); + + let err = BatchPartitioner::new_hash_partitioner(vec![], 0, timer) + .err() + .expect("zero hash partitions should fail") + .to_string(); + + assert!( + err.contains("Hash repartition requires at least one partition"), + "actual: {err}" + ); + } + + #[tokio::test] + async fn one_to_many_round_robin() -> Result<()> { + // define input partitions + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition]; + + // repartition from 1 input to 4 output + let output_partitions = + repartition(&schema, partitions, Partitioning::RoundRobinBatch(4)).await?; + + assert_eq!(4, output_partitions.len()); + for partition in &output_partitions { + assert_eq!(1, partition.len()); + } + assert_eq!(13 * 8, output_partitions[0][0].num_rows()); + assert_eq!(13 * 8, output_partitions[1][0].num_rows()); + assert_eq!(12 * 8, output_partitions[2][0].num_rows()); + assert_eq!(12 * 8, output_partitions[3][0].num_rows()); + + Ok(()) + } + + #[tokio::test] + async fn many_to_one_round_robin() -> Result<()> { + // define input partitions + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + // repartition from 3 input to 1 output + let output_partitions = + repartition(&schema, partitions, Partitioning::RoundRobinBatch(1)).await?; + + assert_eq!(1, output_partitions.len()); + assert_eq!(150 * 8, output_partitions[0][0].num_rows()); + + Ok(()) + } + + #[tokio::test] + async fn many_to_many_round_robin() -> Result<()> { + // define input partitions + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + // repartition from 3 input to 5 output + let output_partitions = + repartition(&schema, partitions, Partitioning::RoundRobinBatch(5)).await?; + + let total_rows_per_partition = 8 * 50 * 3 / 5; + assert_eq!(5, output_partitions.len()); + for partition in output_partitions { + assert_eq!(1, partition.len()); + assert_eq!(total_rows_per_partition, partition[0].num_rows()); + } + + Ok(()) + } + + #[tokio::test] + async fn many_to_many_hash_partition() -> Result<()> { + // define input partitions + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + let output_partitions = repartition( + &schema, + partitions, + Partitioning::Hash(vec![col("c0", &schema)?], 8), + ) + .await?; + + let total_rows: usize = output_partitions + .iter() + .map(|x| x.iter().map(|x| x.num_rows()).sum::<usize>()) + .sum(); + + assert_eq!(8, output_partitions.len()); + assert_eq!(total_rows, 8 * 50 * 3); + + Ok(()) + } + + #[tokio::test] + async fn many_to_many_range_partition() -> Result<()> { + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; + + // create_batch values are [1, 2, 3, 4, 5, 6, 7, 8]; split at 3 and 6 yields + // 2, 3, and 3 rows per batch respectively + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![3, 6])?; + + let output_partitions = repartition(&schema, partitions, partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!(300, partition_row_count(&output_partitions[0])); + assert_eq!(450, partition_row_count(&output_partitions[1])); + assert_eq!(450, partition_row_count(&output_partitions[2])); + assert_eq!( + collect_partition_u32_values(&output_partitions[0]) + .into_iter() + .flatten() + .collect::<HashSet<_>>(), + HashSet::from([1, 2]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[1]) + .into_iter() + .flatten() + .collect::<HashSet<_>>(), + HashSet::from([3, 4, 5]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[2]) + .into_iter() + .flatten() + .collect::<HashSet<_>>(), + HashSet::from([6, 7, 8]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_compound_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![5, 10, 10, 10, 10, 15])), + Arc::new(UInt32Array::from(vec![1, 1, 3, 5, 7, 0])), + ], + )?; + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(1)), + ]), + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(5)), + ]), + ], + )?); + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![(5, 1)], + collect_partition_u32_pairs(&output_partitions[0]) + ); + assert_eq!( + vec![(10, 1), (10, 3)], + collect_partition_u32_pairs(&output_partitions[1]) + ); + assert_eq!( + vec![(10, 5), (10, 7), (15, 0)], + collect_partition_u32_pairs(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_last() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, false), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![None, Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_first() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, true), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![None, Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_asc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![10, 20])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_desc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 20, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(true, false), vec![20, 10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(15), Some(20)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(5), Some(10)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_string_rows() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["bar", "baz", "foo", "qux"])) as ArrayRef, + )])?; + + let schema = batch.schema(); + let expr = col("my_awesome_field", &schema)?; + let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema)); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new_default(expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "foo".to_string(), + ))])], + )?); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?; + + let mut partition_0 = Vec::new(); + let mut stream = exec.execute(0, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + partition_0.push(result?); + } + + let mut partition_1 = Vec::new(); + let mut stream = exec.execute(1, task_ctx)?; + while let Some(result) = stream.next().await { + partition_1.push(result?); + } + + assert_eq!( + vec!["bar", "baz"], + collect_partition_string_values(&partition_0) + ); + assert_eq!( + vec!["foo", "qux"], + collect_partition_string_values(&partition_1) + ); + + Ok(()) + } + + #[test] + fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { + // Three columns so the projection both narrows the schema (required for + // swap) and moves the range key from @0 to @1. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; + + let swapped = repartition + .try_swapping_with_projection(&projection)? + .expect("swap should succeed when projection keeps the range key"); + let swapped_repartition = swapped + .downcast_ref::<RepartitionExec>() + .expect("top node should be RepartitionExec"); + + assert!(swapped_repartition.input().is::<ProjectionExec>()); + let range = expect_range_partitioning(swapped_repartition.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); + assert_eq!( + range.split_points(), + &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])] + ); + + Ok(()) + } + + #[test] + fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { + // Drop a simple range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + // Drop part of a compound range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_when_maintains_order() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source with preserve_order: Range maintains input order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new( + RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )? + .with_preserve_order(), + ); + assert!(repartition.maintains_input_order()[0]); + + match repartition.try_pushdown_sort(ordering.as_ref())? { + SortOrderPushdownResult::Exact { inner } => { + let pushed = inner + .downcast_ref::<RepartitionExec>() + .expect("pushdown should keep RepartitionExec"); + + assert!(pushed.preserve_order()); + assert!(pushed.maintains_input_order()[0]); + + let range = expect_range_partitioning(pushed.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + assert_eq!( + inner.properties().output_ordering().map(|o| o.to_string()), + Some(ordering.to_string()), + "pushed repartition output ordering should match the requested sort" + ); + } + other => panic!("expected Exact sort pushdown, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance() + -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source without preserve_order: Range does not maintain order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new(RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + assert!(!repartition.maintains_input_order()[0]); + + assert!(matches!( + repartition.try_pushdown_sort(ordering.as_ref())?, + SortOrderPushdownResult::Unsupported + )); + + Ok(()) + } + + fn range_partitioning_on_columns( + schema: &SchemaRef, + key_columns: &[&str], + split_points: Vec<Vec<u32>>, + ) -> Result<Partitioning> { + let Some(ordering) = LexOrdering::new( + key_columns + .iter() + .map(|name| { + Ok(PhysicalSortExpr::new( + col(name, schema)?, + SortOptions::default(), + )) + }) + .collect::<Result<Vec<_>>>()?, + ) else { + return exec_err!("range ordering must not be empty"); + }; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points + .into_iter() + .map(|values| { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::UInt32(Some(value))) + .collect(), + ) + }) + .collect(), + )?)) + } + + fn projection_on_columns( + input: &Arc<dyn ExecutionPlan>, + names: &[&str], + ) -> Result<ProjectionExec> { + let exprs = names + .iter() + .map(|name| { + Ok(ProjectionExpr { + expr: col(name, &input.schema())?, + alias: (*name).to_string(), + }) + }) + .collect::<Result<Vec<_>>>()?; + ProjectionExec::try_new(exprs, Arc::clone(input)) + } + + fn expect_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { + match partitioning { + Partitioning::Range(range) => range, + other => panic!("expected Range partitioning, got {other:?}"), + } + } + + /// Test source that claims Exact support for any sort pushdown request. + #[derive(Debug, Clone)] + struct ExactSortPushdownExec { + cache: Arc<PlanProperties>, + } + + impl ExactSortPushdownExec { + fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self { + use crate::execution_plan::{Boundedness, EmissionType}; + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, [ordering]), + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } + } + + impl DisplayAs for ExactSortPushdownExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExactSortPushdownExec") + } + } + + impl ExecutionPlan for ExactSortPushdownExec { + fn name(&self) -> &str { + "ExactSortPushdownExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(EmptyRecordBatchStream::new(self.schema()))) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result<SortOrderPushdownResult<Arc<dyn ExecutionPlan>>> { + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(self.clone()), + }) + } + } + + #[tokio::test] + async fn test_repartition_with_coalescing() -> Result<()> { + let schema = test_schema(false); + // create 50 batches, each having 8 rows + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone()]; + let partitioning = Partitioning::RoundRobinBatch(1); + + let session_config = SessionConfig::new().with_batch_size(200); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + assert_eq!(200, batch.num_rows()); + } + } + Ok(()) + } + + #[tokio::test] + async fn unbounded_input_emits_before_batch_size() -> Result<()> { + let schema = test_schema(false); + let batch = create_batch(); + let source = Arc::new(StreamingTableExec::try_new( + Arc::clone(&schema), + vec![Arc::new(UnboundedTestPartition { + schema: Arc::clone(&schema), + batch: batch.clone(), + })], + None, + vec![], + true, + None, + )?); + let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(1))?; + let session_config = SessionConfig::new().with_batch_size(batch.num_rows() * 2); + let task_ctx = + Arc::new(TaskContext::default().with_session_config(session_config)); + + let mut stream = exec.execute(0, task_ctx)?; + let output = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("unbounded repartition withheld a partial batch") + .expect("unbounded input ended unexpectedly")?; + + assert_eq!(batch, output); + Ok(()) + } + + fn test_schema(nullable: bool) -> Arc<Schema> { + Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::UInt32, + nullable, + )])) + } + + fn u32_range_partitioning( + schema: &SchemaRef, + sort_options: SortOptions, + split_values: Vec<u32>, + ) -> Result<Partitioning> { + let expr = col("c0", schema)?; + Ok(Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(expr, sort_options)].into(), + split_values + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect(), + )?)) + } + + fn partition_row_count(batches: &[RecordBatch]) -> usize { + batches.iter().map(|batch| batch.num_rows()).sum() + } + + fn collect_partition_u32_values(batches: &[RecordBatch]) -> Vec<Option<u32>> { + batches + .iter() + .flat_map(|batch| { + let array = + as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + (0..array.len()) + .map(|idx| { + if array.is_null(idx) { + None + } else { + Some(array.value(idx)) + } + }) + .collect::<Vec<_>>() + }) + .collect() + } + + fn collect_partition_u32_pairs(batches: &[RecordBatch]) -> Vec<(u32, u32)> { + batches + .iter() + .flat_map(|batch| { + let a = as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + let b = as_uint32_array(batch.column(1)).expect("expected UInt32 column"); + (0..a.len()) + .map(|idx| (a.value(idx), b.value(idx))) + .collect::<Vec<_>>() + }) + .collect() + } + + fn collect_partition_string_values(batches: &[RecordBatch]) -> Vec<&str> { + batches + .iter() + .flat_map(|batch| { + let array = + as_string_array(batch.column(0)).expect("expected Utf8 column"); + (0..array.len()) + .map(|idx| array.value(idx)) + .collect::<Vec<_>>() + }) + .collect() + } + + async fn repartition( + schema: &SchemaRef, + input_partitions: Vec<Vec<RecordBatch>>, + partitioning: Partitioning, + ) -> Result<Vec<Vec<RecordBatch>>> { + let task_ctx = Arc::new(TaskContext::default()); + // create physical plan + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + // execute and collect results + let mut output_partitions = vec![]; + for i in 0..exec.partitioning().partition_count() { + // execute this *output* partition and collect all batches + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + let mut batches = vec![]; + while let Some(result) = stream.next().await { + batches.push(result?); + } + output_partitions.push(batches); + } + Ok(output_partitions) + } + + #[tokio::test] + async fn many_to_many_round_robin_within_tokio_task() -> Result<()> { + let handle: SpawnedTask<Result<Vec<Vec<RecordBatch>>>> = + SpawnedTask::spawn(async move { + // define input partitions + let schema = test_schema(false); + let partition = create_vec_batches(50); + let partitions = + vec![partition.clone(), partition.clone(), partition.clone()]; + + // repartition from 3 input to 5 output + repartition(&schema, partitions, Partitioning::RoundRobinBatch(5)).await + }); + + let output_partitions = handle.join().await.unwrap().unwrap(); + + let total_rows_per_partition = 8 * 50 * 3 / 5; + assert_eq!(5, output_partitions.len()); + for partition in output_partitions { + assert_eq!(1, partition.len()); + assert_eq!(total_rows_per_partition, partition[0].num_rows()); + } + + Ok(()) + } + + #[tokio::test] + async fn unsupported_partitioning() { + let task_ctx = Arc::new(TaskContext::default()); + // have to send at least one batch through to provoke error + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef, + )]) + .unwrap(); + + let schema = batch.schema(); + let input = MockExec::new(vec![Ok(batch)], schema); + // This generates an error (partitioning type not supported) + // but only after the plan is executed. The error should be + // returned and no results produced + let partitioning = Partitioning::UnknownPartitioning(1); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap(); + let output_stream = exec.execute(0, task_ctx).unwrap(); + + // Expect that an error is returned + let result_string = crate::common::collect(output_stream) + .await + .unwrap_err() + .to_string(); + assert!( + result_string + .contains("Unsupported repartitioning scheme UnknownPartitioning(1)"), + "actual: {result_string}" + ); + } + + #[tokio::test] + async fn error_for_input_exec() { + // This generates an error on a call to execute. The error + // should be returned and no results produced. + + let task_ctx = Arc::new(TaskContext::default()); + let input = ErrorExec::new(); + let partitioning = Partitioning::RoundRobinBatch(1); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap(); + + // Expect that an error is returned + let result_string = exec.execute(0, task_ctx).err().unwrap().to_string(); + + assert!( + result_string.contains("ErrorExec, unsurprisingly, errored in partition 0"), + "actual: {result_string}" + ); + } + + #[tokio::test] + async fn repartition_with_error_in_stream() { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef, + )]) + .unwrap(); + + // input stream returns one good batch and then one error. The + // error should be returned. + let err = exec_err!("bad data error"); + + let schema = batch.schema(); + let input = MockExec::new(vec![Ok(batch), err], schema); + let partitioning = Partitioning::RoundRobinBatch(1); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap(); + + // Note: this should pass (the stream can be created) but the + // error when the input is executed should get passed back + let output_stream = exec.execute(0, task_ctx).unwrap(); + + // Expect that an error is returned + let result_string = crate::common::collect(output_stream) + .await + .unwrap_err() + .to_string(); + assert!( + result_string.contains("bad data error"), + "actual: {result_string}" + ); + } + + #[tokio::test] + async fn repartition_with_delayed_stream() { + let task_ctx = Arc::new(TaskContext::default()); + let batch1 = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef, + )]) + .unwrap(); + + let batch2 = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["frob", "baz"])) as ArrayRef, + )]) + .unwrap(); + + // The mock exec doesn't return immediately (instead it + // requires the input to wait at least once) + let schema = batch1.schema(); + let expected_batches = vec![batch1.clone(), batch2.clone()]; + let input = MockExec::new(vec![Ok(batch1), Ok(batch2)], schema); + let partitioning = Partitioning::RoundRobinBatch(1); + + let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap(); + + assert_snapshot!(batches_to_sort_string(&expected_batches), @r" + +------------------+ + | my_awesome_field | + +------------------+ + | bar | + | baz | + | foo | + | frob | + +------------------+ + "); + + let output_stream = exec.execute(0, task_ctx).unwrap(); + let batches = crate::common::collect(output_stream).await.unwrap(); + + assert_snapshot!(batches_to_sort_string(&batches), @r" + +------------------+ + | my_awesome_field | + +------------------+ + | bar | + | baz | + | foo | + | frob | + +------------------+ + "); + } + + #[tokio::test] + async fn robin_repartition_with_dropping_output_stream() { + let task_ctx = Arc::new(TaskContext::default()); + let partitioning = Partitioning::RoundRobinBatch(2); + // The barrier exec waits to be pinged + // requires the input to wait at least once) + let input = Arc::new(make_barrier_exec()); + + // partition into two output streams + let exec = RepartitionExec::try_new( + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + partitioning, + ) + .unwrap(); + + let output_stream0 = exec.execute(0, Arc::clone(&task_ctx)).unwrap(); + let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap(); + + // now, purposely drop output stream 0 + // *before* any outputs are produced + drop(output_stream0); + + // Now, start sending input + let mut background_task = JoinSet::new(); + background_task.spawn(async move { + input.wait().await; + }); + + // output stream 1 should *not* error and have one of the input batches + let batches = crate::common::collect(output_stream1).await.unwrap(); + + assert_snapshot!(batches_to_sort_string(&batches), @r" + +------------------+ + | my_awesome_field | + +------------------+ + | baz | + | frob | + | gar | + | goo | + +------------------+ + "); + } + + #[tokio::test] + // As the hash results might be different on different platforms or + // with different compilers, we will compare the same execution with + // and without dropping the output stream. + async fn hash_repartition_with_dropping_output_stream() { + let task_ctx = Arc::new(TaskContext::default()); + let partitioning = Partitioning::Hash( + vec![Arc::new(crate::expressions::Column::new( + "my_awesome_field", + 0, + ))], + 2, + ); + + // We first collect the results without dropping the output stream. + let input = Arc::new(make_barrier_exec()); + let exec = RepartitionExec::try_new( + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + partitioning.clone(), + ) + .unwrap(); + let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap(); + let mut background_task = JoinSet::new(); + background_task.spawn(async move { + input.wait().await; + }); + let batches_without_drop = crate::common::collect(output_stream1).await.unwrap(); + + // run some checks on the result + let items_vec = str_batches_to_vec(&batches_without_drop); + let items_set: HashSet<&str> = items_vec.iter().copied().collect(); + assert_eq!(items_vec.len(), items_set.len()); + let source_str_set: HashSet<&str> = + ["foo", "bar", "frob", "baz", "goo", "gar", "grob", "gaz"] + .iter() + .copied() + .collect(); + assert_eq!(items_set.difference(&source_str_set).count(), 0); + + // Now do the same but dropping the stream before waiting for the barrier + let input = Arc::new(make_barrier_exec()); + let exec = RepartitionExec::try_new( + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + partitioning, + ) + .unwrap(); + let output_stream0 = exec.execute(0, Arc::clone(&task_ctx)).unwrap(); + let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap(); + // now, purposely drop output stream 0 + // *before* any outputs are produced + drop(output_stream0); + let mut background_task = JoinSet::new(); + background_task.spawn(async move { + input.wait().await; + }); + let batches_with_drop = crate::common::collect(output_stream1).await.unwrap(); + + let items_vec_with_drop = str_batches_to_vec(&batches_with_drop); + let items_set_with_drop: HashSet<&str> = + items_vec_with_drop.iter().copied().collect(); + assert_eq!( + items_set_with_drop.symmetric_difference(&items_set).count(), + 0 + ); + } + + fn str_batches_to_vec(batches: &[RecordBatch]) -> Vec<&str> { + batches + .iter() + .flat_map(|batch| { + assert_eq!(batch.columns().len(), 1); + let string_array = as_string_array(batch.column(0)) + .expect("Unexpected type for repartitioned batch"); + + string_array + .iter() + .map(|v| v.expect("Unexpected null")) + .collect::<Vec<_>>() + }) + .collect::<Vec<_>>() + } + + /// Create a BarrierExec that returns two partitions of two batches each + fn make_barrier_exec() -> BarrierExec { + let batch1 = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["foo", "bar"])) as ArrayRef, + )]) + .unwrap(); + + let batch2 = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["frob", "baz"])) as ArrayRef, + )]) + .unwrap(); + + let batch3 = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["goo", "gar"])) as ArrayRef, + )]) + .unwrap(); + + let batch4 = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["grob", "gaz"])) as ArrayRef, + )]) + .unwrap(); + + // The barrier exec waits to be pinged + // requires the input to wait at least once) + let schema = batch1.schema(); + BarrierExec::new(vec![vec![batch1, batch2], vec![batch3, batch4]], schema) + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 2)); + let refs = blocking_exec.refs(); + let repartition_exec = Arc::new(RepartitionExec::try_new( + blocking_exec, + Partitioning::UnknownPartitioning(1), + )?); + + let fut = collect(repartition_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn hash_repartition_avoid_empty_batch() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(StringArray::from(vec!["foo"])) as ArrayRef, + )]) + .unwrap(); + let partitioning = Partitioning::Hash( + vec![Arc::new(crate::expressions::Column::new("a", 0))], + 2, + ); + let schema = batch.schema(); + let input = MockExec::new(vec![Ok(batch)], schema); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning).unwrap(); + let output_stream0 = exec.execute(0, Arc::clone(&task_ctx)).unwrap(); + let batch0 = crate::common::collect(output_stream0).await.unwrap(); + let output_stream1 = exec.execute(1, Arc::clone(&task_ctx)).unwrap(); + let batch1 = crate::common::collect(output_stream1).await.unwrap(); + assert!(batch0.is_empty() || batch1.is_empty()); + Ok(()) + } + + #[tokio::test] + async fn repartition_with_spilling() -> Result<()> { + // Test that repartition successfully spills to disk when memory is constrained + let schema = test_schema(false); + let partition = create_vec_batches(50); + let input_partitions = vec![partition]; + let partitioning = Partitioning::RoundRobinBatch(4); + + // Set up context with very tight memory limit to force spilling + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1, 1.0) + .build_arc()?; + + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + // Collect all partitions - should succeed by spilling to disk + let mut total_rows = 0; + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + total_rows += batch.num_rows(); + } + } + + // Verify we got all the data (50 batches * 8 rows each) + assert_eq!(total_rows, 50 * 8); + + // Verify spilling metrics to confirm spilling actually happened + let metrics = exec.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spill_count > 0, but got {:?}", + metrics.spill_count() + ); + println!("Spilled {} times", metrics.spill_count().unwrap()); + assert!( + metrics.spilled_bytes().unwrap() > 0, + "Expected spilled_bytes > 0, but got {:?}", + metrics.spilled_bytes() + ); + println!( + "Spilled {} bytes in {} spills", + metrics.spilled_bytes().unwrap(), + metrics.spill_count().unwrap() + ); + assert!( + metrics.spilled_rows().unwrap() > 0, + "Expected spilled_rows > 0, but got {:?}", + metrics.spilled_rows() + ); + println!("Spilled {} rows", metrics.spilled_rows().unwrap()); + + Ok(()) + } + + #[tokio::test] + async fn repartition_with_partial_spilling() -> Result<()> { + // Test that repartition can handle partial spilling (some batches in memory, some spilled) + let schema = test_schema(false); + let partition = create_vec_batches(50); + let input_partitions = vec![partition]; + let partitioning = Partitioning::RoundRobinBatch(4); + + // With `batch_size = 1024` and a single UInt32 column, each + // coalesced residual is ~4 KiB. An 8 KiB pool fits one and forces + // the rest to spill. + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(8 * 1024, 1.0) + .build_arc()?; + + let session_config = SessionConfig::new().with_batch_size(1024); + let task_ctx = TaskContext::default() + .with_runtime(runtime) + .with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + // Collect all partitions - should succeed with partial spilling + let mut total_rows = 0; + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + total_rows += batch.num_rows(); + } + } + + // Verify we got all the data (50 batches * 8 rows each) + assert_eq!(total_rows, 50 * 8); + + // Verify partial spilling metrics + let metrics = exec.metrics().unwrap(); + let spill_count = metrics.spill_count().unwrap(); + let spilled_rows = metrics.spilled_rows().unwrap(); + let spilled_bytes = metrics.spilled_bytes().unwrap(); + + assert!( + spill_count > 0, + "Expected some spilling to occur, but got spill_count={spill_count}" + ); + assert!( + spilled_rows > 0 && spilled_rows < total_rows, + "Expected partial spilling (0 < spilled_rows < {total_rows}), but got spilled_rows={spilled_rows}" + ); + assert!( + spilled_bytes > 0, + "Expected some bytes to be spilled, but got spilled_bytes={spilled_bytes}" + ); + + println!( + "Partial spilling: spilled {} out of {} rows ({:.1}%) in {} spills, {} bytes", + spilled_rows, + total_rows, + (spilled_rows as f64 / total_rows as f64) * 100.0, + spill_count, + spilled_bytes + ); + + Ok(()) + } + + #[tokio::test] + async fn repartition_without_spilling() -> Result<()> { + // Test that repartition does not spill when there's ample memory + let schema = test_schema(false); + let partition = create_vec_batches(50); + let input_partitions = vec![partition]; + let partitioning = Partitioning::RoundRobinBatch(4); + + // Set up context with generous memory limit - no spilling should occur + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(10 * 1024 * 1024, 1.0) // 10MB + .build_arc()?; + + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + // Collect all partitions - should succeed without spilling + let mut total_rows = 0; + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + total_rows += batch.num_rows(); + } + } + + // Verify we got all the data (50 batches * 8 rows each) + assert_eq!(total_rows, 50 * 8); + + // Verify no spilling occurred + let metrics = exec.metrics().unwrap(); + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spilling, but got spill_count={:?}", + metrics.spill_count() + ); + assert_eq!( + metrics.spilled_bytes(), + Some(0), + "Expected no bytes spilled, but got spilled_bytes={:?}", + metrics.spilled_bytes() + ); + assert_eq!( + metrics.spilled_rows(), + Some(0), + "Expected no rows spilled, but got spilled_rows={:?}", + metrics.spilled_rows() + ); + + println!("No spilling occurred - all data processed in memory"); + + Ok(()) + } + + #[tokio::test] + async fn oom() -> Result<()> { + use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; + + // Test that repartition fails with OOM when disk manager is disabled + let schema = test_schema(false); + let partition = create_vec_batches(50); + let input_partitions = vec![partition]; + let partitioning = Partitioning::RoundRobinBatch(4); + + // Setup context with memory limit but NO disk manager (explicitly disabled) + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::Disabled), + ) + .build_arc()?; + + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + // Attempt to execute - should fail with ResourcesExhausted error + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + let err = stream.next().await.unwrap().unwrap_err(); + let err = err.find_root(); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "Wrong error type: {err}", + ); + } + + Ok(()) + } + + /// Create vector batches + fn create_vec_batches(n: usize) -> Vec<RecordBatch> { + let batch = create_batch(); + std::iter::repeat_n(batch, n).collect() + } + + /// Create batch + fn create_batch() -> RecordBatch { + let schema = test_schema(false); + RecordBatch::try_new( + schema, + vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))], + ) + .unwrap() + } + + /// Create batches with sequential values for ordering tests + fn create_ordered_batches(num_batches: usize) -> Vec<RecordBatch> { + let schema = test_schema(false); + (0..num_batches) + .map(|i| { + let start = (i * 8) as u32; + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from( + (start..start + 8).collect::<Vec<_>>(), + ))], + ) + .unwrap() + }) + .collect() + } + + #[tokio::test] + async fn test_repartition_ordering_with_spilling() -> Result<()> { + // Test that repartition preserves ordering when spilling occurs + // This tests the state machine fix where we must block on spill_stream + // when a Spilled marker is received, rather than continuing to poll the channel + + let schema = test_schema(false); + // Create batches with sequential values: batch 0 has [0,1,2,3,4,5,6,7], + // batch 1 has [8,9,10,11,12,13,14,15], etc. + let partition = create_ordered_batches(20); + let input_partitions = vec![partition]; + + // Use RoundRobinBatch to ensure predictable ordering + let partitioning = Partitioning::RoundRobinBatch(2); + + // Set up context with very tight memory limit to force spilling + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1, 1.0) + .build_arc()?; + + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + // Collect all output partitions + let mut all_batches = Vec::new(); + for i in 0..exec.partitioning().partition_count() { + let mut partition_batches = Vec::new(); + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + partition_batches.push(batch); + } + all_batches.push(partition_batches); + } + + // Verify spilling occurred + let metrics = exec.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur, but spill_count = 0" + ); + + // Verify ordering is preserved within each partition + // With RoundRobinBatch, even batches go to partition 0, odd batches to partition 1 + for (partition_idx, batches) in all_batches.iter().enumerate() { + let mut last_value = None; + for batch in batches { + let array = batch + .column(0) + .as_any() + .downcast_ref::<UInt32Array>() + .unwrap(); + + for i in 0..array.len() { + let value = array.value(i); + if let Some(last) = last_value { + assert!( + value > last, + "Ordering violated in partition {partition_idx}: {value} is not greater than {last}" + ); + } + last_value = Some(value); + } + } + } + + Ok(()) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::TestMemoryExec; + use crate::union::UnionExec; + use arrow::array::{UInt32Array, record_batch}; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::assert_batches_eq; + use datafusion_common::config::ConfigNonZeroUsize; + + use datafusion_physical_expr::expressions::col; + + /// Asserts that the plan is as expected + /// + /// `$EXPECTED_PLAN_LINES`: input plan + /// `$PLAN`: the plan to optimized + macro_rules! assert_plan { + ($PLAN: expr, @ $EXPECTED: expr) => { + let formatted = crate::displayable($PLAN).indent(true).to_string(); + + insta::assert_snapshot!( + formatted, + @$EXPECTED + ); + }; + } + + #[tokio::test] + async fn test_preserve_order() -> Result<()> { + let schema = test_schema(); + let sort_exprs = sort_exprs(&schema); + let source1 = sorted_memory_exec(&schema, sort_exprs.clone()); + let source2 = sorted_memory_exec(&schema, sort_exprs); + // output has multiple partitions, and is sorted + let union = UnionExec::try_new(vec![source1, source2])?; + let exec = RepartitionExec::try_new(union, Partitioning::RoundRobinBatch(10))? + .with_preserve_order(); + + // Repartition should preserve order + assert_plan!(&exec, @r" + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2, preserve_order=true, sort_exprs=c0@0 ASC + UnionExec + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + "); + Ok(()) + } + + #[tokio::test] + async fn test_preserve_order_one_partition() -> Result<()> { + let schema = test_schema(); + let sort_exprs = sort_exprs(&schema); + let source = sorted_memory_exec(&schema, sort_exprs); + // output is sorted, but has only a single partition, so no need to sort + let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(10))? + .with_preserve_order(); + + // Repartition should not preserve order + assert_plan!(&exec, @r" + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + "); + + Ok(()) + } + + #[tokio::test] + async fn test_preserve_order_input_not_sorted() -> Result<()> { + let schema = test_schema(); + let source1 = memory_exec(&schema); + let source2 = memory_exec(&schema); + // output has multiple partitions, but is not sorted + let union = UnionExec::try_new(vec![source1, source2])?; + let exec = RepartitionExec::try_new(union, Partitioning::RoundRobinBatch(10))? + .with_preserve_order(); + + // Repartition should not preserve order, as there is no order to preserve + assert_plan!(&exec, @r" + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 + UnionExec + DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] + "); + Ok(()) + } + + #[tokio::test] + async fn test_preserve_order_with_spilling() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + // Create sorted input data across multiple partitions + // Partition1: [1,3], [5,7], [9,11] + // Partition2: [2,4], [6,8], [10,12] + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap(); + let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap(); + let batch5 = record_batch!(("c0", UInt32, [9, 11])).unwrap(); + let batch6 = record_batch!(("c0", UInt32, [10, 12])).unwrap(); + let schema = batch1.schema(); + let sort_exprs = LexOrdering::new([PhysicalSortExpr { + expr: col("c0", &schema).unwrap(), + options: SortOptions::default().asc(), + }]) + .unwrap(); + let partition1 = vec![batch1.clone(), batch3.clone(), batch5.clone()]; + let partition2 = vec![batch2.clone(), batch4.clone(), batch6.clone()]; + let input_partitions = vec![partition1, partition2]; + + // Set up context with tight memory limit to force spilling + // Sorting needs some non-spillable memory, so 608 bytes should force spilling while still allowing the query to complete + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(608, 1.0) + .build_arc()?; + + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // Create physical plan with order preservation + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?; + let exec = Arc::new(exec); + let exec = Arc::new(TestMemoryExec::update_cache(&exec)); + // Repartition into 3 partitions with order preservation + // We expect 1 batch per output partition after repartitioning + let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))? + .with_preserve_order(); + + let mut batches = vec![]; + + // Collect all partitions - should succeed by spilling to disk + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + batches.push(batch); + } + } + + #[rustfmt::skip] + let expected = [ + [ + "+----+", + "| c0 |", + "+----+", + "| 1 |", + "| 2 |", + "| 3 |", + "| 4 |", + "+----+", + ], + [ + "+----+", + "| c0 |", + "+----+", + "| 5 |", + "| 6 |", + "| 7 |", + "| 8 |", + "+----+", + ], + [ + "+----+", + "| c0 |", + "+----+", + "| 9 |", + "| 10 |", + "| 11 |", + "| 12 |", + "+----+", + ], + ]; + + for (batch, expected) in batches.iter().zip(expected.iter()) { + assert_batches_eq!(expected, std::slice::from_ref(batch)); + } + + // We should have spilled + let metrics = exec.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur for order-preserving repartition at this \ + memory limit. If this fails, the memory limit may need adjustment." + ); + Ok(()) + } + + /// Regression test for order preservation across spill *file rotation*. + /// + /// A `preserve_order` repartition relies on each per-(input, output) spill pool delivering + /// batches in strict FIFO order (see [`spill_pool::spsc_channel`] / [`SpillPoolSink`]). This uses + /// the same memory profile as [`Self::test_preserve_order_with_spilling`] — which is tuned to + /// force spilling while still completing — but additionally sets `max_spill_file_size_bytes` + /// to 1 so every spilled batch lands in its own file. That exercises the FIFO-across-rotation + /// path: if ordering were lost across rotated files (e.g. by feeding an ordered pool with a + /// shared multi-producer writer), the downstream `StreamingMerge` would emit out-of-order rows + /// and the sortedness assertion below would fail. + #[tokio::test] + async fn test_preserve_order_with_spill_file_rotation() -> Result<()> { + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + // Same sorted input as `test_preserve_order_with_spilling`: + // Partition1: [1,3], [5,7], [9,11]; Partition2: [2,4], [6,8], [10,12] + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap(); + let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap(); + let batch5 = record_batch!(("c0", UInt32, [9, 11])).unwrap(); + let batch6 = record_batch!(("c0", UInt32, [10, 12])).unwrap(); + let schema = batch1.schema(); + let sort_exprs = LexOrdering::new([PhysicalSortExpr { + expr: col("c0", &schema).unwrap(), + options: SortOptions::default().asc(), + }]) + .unwrap(); + let partition1 = vec![batch1, batch3, batch5]; + let partition2 = vec![batch2, batch4, batch6]; + let input_partitions = vec![partition1, partition2]; + + // Force a new spill file per spilled batch to exercise FIFO across rotation. + let mut session_config = SessionConfig::new(); + session_config + .options_mut() + .execution + .max_spill_file_size_bytes = ConfigNonZeroUsize::try_new(1).unwrap(); + // Same tight limit as `test_preserve_order_with_spilling`: forces spilling while leaving + // the merge enough non-spillable headroom to complete. + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(608, 1.0) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))? + .with_preserve_order(); + + // Each output partition merges sorted substreams, so its rows must be non-decreasing. + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + let mut last: Option<u32> = None; + while let Some(result) = stream.next().await { + let batch = result?; + let col = batch + .column(0) + .as_any() + .downcast_ref::<UInt32Array>() + .unwrap(); + for r in 0..col.len() { + let v = col.value(r); + if let Some(prev) = last { + assert!( + prev <= v, + "output partition {i} not sorted: {prev} came before {v}" + ); + } + last = Some(v); + } + } + } + + let metrics = exec.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur for order-preserving repartition at this \ + memory limit. If this fails, the memory limit may need adjustment." + ); + Ok(()) + } + + #[tokio::test] + async fn test_hash_partitioning_with_spilling() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + // Create input data similar to the round-robin test + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap(); + let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap(); + let schema = batch1.schema(); + + let partition1 = vec![batch1.clone(), batch3.clone()]; + let partition2 = vec![batch2.clone(), batch4.clone()]; + let input_partitions = vec![partition1, partition2]; + + // Set up context with memory limit to test hash partitioning with spilling infrastructure + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1, 1.0) + .build_arc()?; + + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + // Create physical plan with hash partitioning + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(exec); + let exec = Arc::new(TestMemoryExec::update_cache(&exec)); + // Hash partition into 2 partitions by column c0 + let hash_expr = col("c0", &schema)?; + let exec = + RepartitionExec::try_new(exec, Partitioning::Hash(vec![hash_expr], 2))?; + + // Collect all partitions concurrently using JoinSet - this prevents deadlock + // where the distribution channel gate closes when all output channels are full + let mut join_set = tokio::task::JoinSet::new(); + for i in 0..exec.partitioning().partition_count() { + let stream = exec.execute(i, Arc::clone(&task_ctx))?; + join_set.spawn(async move { + let mut count = 0; + futures::pin_mut!(stream); + while let Some(result) = stream.next().await { + let batch = result?; + count += batch.num_rows(); + } + Ok::<usize, DataFusionError>(count) + }); + } + + // Wait for all partitions and sum the rows + let mut total_rows = 0; + while let Some(result) = join_set.join_next().await { + total_rows += result.unwrap()?; + } + + // Verify we got all rows back + let all_batches = [batch1, batch2, batch3, batch4]; + let expected_rows: usize = all_batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, expected_rows); + + // Verify metrics are available + let metrics = exec.metrics().unwrap(); + // Just verify the metrics can be retrieved (spilling may or may not occur) + let spill_count = metrics.spill_count().unwrap_or(0); + assert!(spill_count > 0); + let spilled_bytes = metrics.spilled_bytes().unwrap_or(0); + assert!(spilled_bytes > 0); + let spilled_rows = metrics.spilled_rows().unwrap_or(0); + assert!(spilled_rows > 0); + + Ok(()) + } + + #[tokio::test] + async fn test_repartition() -> Result<()> { + let schema = test_schema(); + let sort_exprs = sort_exprs(&schema); + let source = sorted_memory_exec(&schema, sort_exprs); + // output is sorted, but has only a single partition, so no need to sort + let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(10))? + .repartitioned(20, &Default::default())? + .unwrap(); + + // Repartition should not preserve order + assert_plan!(exec.as_ref(), @r" + RepartitionExec: partitioning=RoundRobinBatch(20), input_partitions=1, maintains_sort_order=true + DataSourceExec: partitions=1, partition_sizes=[0], output_ordering=c0@0 ASC + "); + Ok(()) + } + + #[test] + fn test_range_repartitioned_returns_none() -> Result<()> { + let schema = test_schema(); + let source = memory_exec(&schema); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new( + col("c0", &schema)?, + SortOptions::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let exec = RepartitionExec::try_new(source, partitioning)?; + + let mut expressions = vec![]; + exec.apply_expressions(&mut |expr| { + expressions.push(expr.to_string()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(expressions, ["c0@0"]); + + // Range partition count is fixed by split points, so repartitioned() + // cannot change it to an arbitrary target. + let result = exec.repartitioned(10, &Default::default())?; + assert!( + result.is_none(), + "range repartitioning should not support changing partition count" + ); + Ok(()) + } + + fn test_schema() -> Arc<Schema> { + Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) + } + + fn sort_exprs(schema: &Schema) -> LexOrdering { + [PhysicalSortExpr { + expr: col("c0", schema).unwrap(), + options: SortOptions::default(), + }] + .into() + } + + fn memory_exec(schema: &SchemaRef) -> Arc<dyn ExecutionPlan> { + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(schema), None).unwrap() + } + + fn sorted_memory_exec( + schema: &SchemaRef, + sort_exprs: LexOrdering, + ) -> Arc<dyn ExecutionPlan> { + let exec = TestMemoryExec::try_new(&[vec![]], Arc::clone(schema), None) + .unwrap() + .try_with_sort_information(vec![sort_exprs]) + .unwrap(); + let exec = Arc::new(exec); + Arc::new(TestMemoryExec::update_cache(&exec)) + } + + /// preserve_order repartition should not double-count + /// output rows. + #[tokio::test] + async fn test_preserve_order_output_rows_not_double_counted() -> Result<()> { + use datafusion_execution::TaskContext; + + // Two sorted input partitions, 2 rows each (4 total) + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let schema = batch1.schema(); + let sort_exprs = sort_exprs(&schema); + + let input_partitions = vec![vec![batch1], vec![batch2]]; + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?; + let exec = Arc::new(exec); + let exec = Arc::new(TestMemoryExec::update_cache(&exec)); + + let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))? + .with_preserve_order(); + + let task_ctx = Arc::new(TaskContext::default()); + let mut total_rows = 0; + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + total_rows += result?.num_rows(); + } + } + + assert_eq!(total_rows, 4, "actual rows collected should be 4"); + + let metrics = exec.metrics().unwrap(); + let reported_output_rows = metrics.output_rows().unwrap(); + assert_eq!( + reported_output_rows, total_rows, + "metrics output_rows ({reported_output_rows}) should match \ + actual rows collected ({total_rows}), not double-count" + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/scalar_subquery.rs b/native/vendor/datafusion-physical-plan/src/scalar_subquery.rs new file mode 100644 index 0000000..f2b7c5e --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/scalar_subquery.rs
@@ -0,0 +1,673 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Execution plan for uncorrelated scalar subqueries. +//! +//! [`ScalarSubqueryExec`] wraps a main input plan and a set of subquery plans. +//! At execution time, it runs each subquery exactly once, extracts the scalar +//! result, and populates a shared [`ScalarSubqueryResults`] container that +//! [`ScalarSubqueryExpr`] instances hold directly and read from by index. +//! +//! [`ScalarSubqueryExpr`]: datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr + +use std::fmt; +use std::sync::Arc; + +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, ScalarValue, Statistics, exec_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_physical_expr::PhysicalExpr; + +use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties}; +use crate::joins::utils::{OnceAsync, OnceFut}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; + +use futures::StreamExt; +use futures::TryStreamExt; + +/// Links a scalar subquery's execution plan to its index in the shared results +/// container. The [`ScalarSubqueryExec`] that owns these links populates +/// `results[index]` at execution time, and [`ScalarSubqueryExpr`] instances +/// with the same index read from it. +/// +/// [`ScalarSubqueryExpr`]: datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr +#[derive(Debug, Clone)] +pub struct ScalarSubqueryLink { + /// The physical plan for the subquery. + pub plan: Arc<dyn ExecutionPlan>, + /// Index into the shared results container. + pub index: SubqueryIndex, +} + +/// Manages execution of uncorrelated scalar subqueries for a single plan +/// level. +/// +/// From a query-results perspective, this node is a pass-through: it yields +/// the same batches as its main input and exists only to populate scalar +/// subquery results as a side effect before those batches are produced. +/// +/// The first child node is the **main input plan**, whose batches are passed +/// through unchanged. The remaining children are **subquery plans**, each of +/// which must produce exactly zero or one row. Before any batches from the main +/// input are yielded, all subquery plans are executed and their scalar results +/// are stored in a shared [`ScalarSubqueryResults`] container owned by this +/// node. [`ScalarSubqueryExpr`] nodes embedded in the main input's expressions +/// hold the same container and read from it by index. +/// +/// All subqueries are evaluated eagerly when the first output partition is +/// requested, before any rows from the main input are produced. +/// +/// TODO: Consider overlapping computation of the subqueries with evaluating the +/// main query. +/// +/// [`ScalarSubqueryExpr`]: datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr +#[derive(Debug)] +pub struct ScalarSubqueryExec { + /// The main input plan whose output is passed through. + input: Arc<dyn ExecutionPlan>, + /// Subquery plans and their result indexes. + subqueries: Vec<ScalarSubqueryLink>, + /// Shared one-time async computation of subquery results. + subquery_future: Arc<OnceAsync<()>>, + /// Shared results container; the corresponding `ScalarSubqueryExpr` + /// nodes in the input plan hold the same underlying container. + results: ScalarSubqueryResults, + /// Cached plan properties (copied from input). + cache: Arc<PlanProperties>, +} + +impl ScalarSubqueryExec { + pub fn new( + input: Arc<dyn ExecutionPlan>, + subqueries: Vec<ScalarSubqueryLink>, + results: ScalarSubqueryResults, + ) -> Self { + let cache = Arc::clone(input.properties()); + Self { + input, + subqueries, + subquery_future: Arc::default(), + results, + cache, + } + } + + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + pub fn subqueries(&self) -> &[ScalarSubqueryLink] { + &self.subqueries + } + + pub fn results(&self) -> &ScalarSubqueryResults { + &self.results + } + + /// Returns a per-child bool vec that is `true` for the main input + /// (child 0) and `false` for every subquery child. + fn true_for_input_only(&self) -> Vec<bool> { + std::iter::once(true) + .chain(std::iter::repeat_n(false, self.subqueries.len())) + .collect() + } +} + +impl DisplayAs for ScalarSubqueryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "ScalarSubqueryExec: subqueries={}", + self.subqueries.len() + ) + } + DisplayFormatType::TreeRender => { + write!(f, "") + } + } + } +} + +impl ExecutionPlan for ScalarSubqueryExec { + fn name(&self) -> &'static str { + "ScalarSubqueryExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + let mut children = vec![&self.input]; + for sq in &self.subqueries { + children.push(&sq.plan); + } + children + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + // First child is the main input, the rest are subquery plans. + let input = children.remove(0); + let subqueries = self + .subqueries + .iter() + .zip(children) + .map(|(sq, new_plan)| ScalarSubqueryLink { + plan: new_plan, + index: sq.index, + }) + .collect(); + Ok(Arc::new(ScalarSubqueryExec::new( + input, + subqueries, + self.results.clone(), + ))) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + self.results.clear(); + Ok(Arc::new(ScalarSubqueryExec { + input: Arc::clone(&self.input), + subqueries: self.subqueries.clone(), + subquery_future: Arc::default(), + results: self.results.clone(), + cache: Arc::clone(&self.cache), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let subqueries = self.subqueries.clone(); + let results = self.results.clone(); + let planning_ctx = Arc::clone(&context); + let mut subquery_future = self.subquery_future.try_once(move || { + Ok(async move { execute_subqueries(subqueries, results, planning_ctx).await }) + })?; + let input = Arc::clone(&self.input); + let schema = self.schema(); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::once(async move { + // Execute all subqueries exactly once, even when multiple + // partitions call execute() concurrently. + wait_for_subqueries(&mut subquery_future).await?; + + // Now that the subqueries have finished execution, we can + // safely execute the main input + input.execute(partition, context) + }) + .try_flatten(), + ))) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn maintains_input_order(&self) -> Vec<bool> { + // Only the main input (first child); subquery children don't contribute + // to ordering. + self.true_for_input_only() + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + // ScalarSubqueryExec is a pass-through coordinator: it does not + // benefit from repartitioning any child directly below it. + vec![false; self.subqueries.len() + 1] + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + // Only `self.input` (child 0) is used; the subqueries are skipped. + let mut requests = vec![ChildStats::Skip; 1 + self.subqueries.len()]; + requests[0] = ChildStats::At(partition); + requests + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::clone(&input_stats[0])) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + let input = ctx.encode_child(self.input())?; + // Subquery indices are positional and recovered during decoding. + let subqueries = + ctx.encode_children(self.subqueries().iter().map(|subquery| &subquery.plan))?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::ScalarSubquery(Box::new( + protobuf::ScalarSubqueryExecNode { + input: Some(Box::new(input)), + subqueries, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ScalarSubqueryExec { + /// Reconstruct a [`ScalarSubqueryExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let scalar_subquery = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::ScalarSubquery, + "ScalarSubqueryExec", + ); + let results = ScalarSubqueryResults::new(scalar_subquery.subqueries.len()); + let input_node = scalar_subquery.input.as_deref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ScalarSubqueryExec is missing required field 'input'" + ) + })?; + // The input's ScalarSubqueryExpr nodes must share this results container. + let input = + ctx.decode_child_with_scalar_subquery_results(input_node, results.clone())?; + let subqueries = scalar_subquery + .subqueries + .iter() + .enumerate() + .map(|(index, plan)| { + Ok(ScalarSubqueryLink { + plan: ctx.decode_child(plan)?, + index: SubqueryIndex::new(index), + }) + }) + .collect::<Result<Vec<_>>>()?; + + Ok(Arc::new(Self::new(input, subqueries, results))) + } +} + +/// Wait for the subquery execution future to complete. +async fn wait_for_subqueries(fut: &mut OnceFut<()>) -> Result<()> { + std::future::poll_fn(|cx| fut.get_shared(cx)).await?; + Ok(()) +} + +async fn execute_subqueries( + subqueries: Vec<ScalarSubqueryLink>, + results: ScalarSubqueryResults, + context: Arc<TaskContext>, +) -> Result<()> { + // Evaluate subqueries in parallel; wait for them all to finish evaluation + // before returning. + let futures = subqueries.iter().map(|sq| { + let plan = Arc::clone(&sq.plan); + let ctx = Arc::clone(&context); + let results = results.clone(); + let index = sq.index; + async move { + let value = execute_scalar_subquery(plan, ctx).await?; + results.set(index, value)?; + Ok(()) as Result<()> + } + }); + futures::future::try_join_all(futures).await?; + Ok(()) +} + +/// Execute a single subquery plan and extract the scalar value. +/// Returns NULL for 0 rows, the scalar value for exactly 1 row, +/// or an error for >1 rows. +async fn execute_scalar_subquery( + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, +) -> Result<ScalarValue> { + let schema = plan.schema(); + if schema.fields().len() != 1 { + // Should be enforced by the physical planner. + return internal_err!( + "Scalar subquery must return exactly one column, got {}", + schema.fields().len() + ); + } + + let mut stream = crate::execute_stream(plan, context)?; + let mut result: Option<ScalarValue> = None; + + while let Some(batch) = stream.next().await.transpose()? { + if batch.num_rows() == 0 { + continue; + } + if result.is_some() || batch.num_rows() > 1 { + return exec_err!("Scalar subquery returned more than one row"); + } + result = Some(ScalarValue::try_from_array(batch.column(0), 0)?); + } + + // 0 rows → typed NULL per SQL semantics + match result { + Some(v) => Ok(v), + None => ScalarValue::try_from(schema.field(0).data_type()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test::{self, TestMemoryExec}; + use crate::{ + execution_plan::reset_plan_states, + projection::{ProjectionExec, ProjectionExpr}, + }; + + use std::sync::atomic::{AtomicUsize, Ordering}; + + use crate::test::exec::ErrorExec; + use arrow::array::{Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + + enum ExpectedSubqueryResult { + Value(ScalarValue), + Error(&'static str), + } + + #[derive(Debug)] + struct CountingExec { + inner: Arc<dyn ExecutionPlan>, + execute_calls: Arc<AtomicUsize>, + } + + impl CountingExec { + fn new(inner: Arc<dyn ExecutionPlan>, execute_calls: Arc<AtomicUsize>) -> Self { + Self { + inner, + execute_calls, + } + } + } + + impl DisplayAs for CountingExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CountingExec") + } + DisplayFormatType::TreeRender => write!(f, ""), + } + } + } + + impl ExecutionPlan for CountingExec { + fn name(&self) -> &'static str { + "CountingExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.inner] + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(Self::new( + children.remove(0), + Arc::clone(&self.execute_calls), + ))) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + self.execute_calls.fetch_add(1, Ordering::SeqCst); + self.inner.execute(partition, context) + } + } + + fn make_subquery_plan(batches: Vec<RecordBatch>) -> Arc<dyn ExecutionPlan> { + let schema = batches[0].schema(); + TestMemoryExec::try_new_exec(&[batches], schema, None).unwrap() + } + + fn int32_batch(values: Vec<i32>) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(values))]).unwrap() + } + + fn empty_int64_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![] as Vec<i64>))]) + .unwrap() + } + + fn placeholder_input() -> Arc<dyn ExecutionPlan> { + Arc::new(crate::placeholder_row::PlaceholderRowExec::new( + test::aggr_test_schema(), + )) + } + + fn single_subquery_exec( + input: Arc<dyn ExecutionPlan>, + subquery_plan: Arc<dyn ExecutionPlan>, + results: ScalarSubqueryResults, + ) -> ScalarSubqueryExec { + ScalarSubqueryExec::new( + input, + vec![ScalarSubqueryLink { + plan: subquery_plan, + index: SubqueryIndex::new(0), + }], + results, + ) + } + + fn scalar_subquery_projection_input( + results: ScalarSubqueryResults, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(ScalarSubqueryExpr::new( + DataType::Int32, + false, + SubqueryIndex::new(0), + results, + )), + alias: "sq".to_string(), + }], + placeholder_input(), + )?)) + } + + fn extract_single_int32_value(batches: &[RecordBatch]) -> i32 { + assert_eq!(batches.len(), 1); + let values = batches[0] + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(values.len(), 1); + values.value(0) + } + + #[tokio::test] + async fn test_execute_scalar_subquery_row_count_semantics() -> Result<()> { + for (name, plan, expected) in [ + ( + "single_row", + make_subquery_plan(vec![int32_batch(vec![42])]), + ExpectedSubqueryResult::Value(ScalarValue::Int32(Some(42))), + ), + ( + "zero_rows", + make_subquery_plan(vec![empty_int64_batch()]), + ExpectedSubqueryResult::Value(ScalarValue::Int64(None)), + ), + ( + "multiple_rows", + make_subquery_plan(vec![int32_batch(vec![1, 2, 3])]), + ExpectedSubqueryResult::Error("more than one row"), + ), + ] { + let actual = + execute_scalar_subquery(plan, Arc::new(TaskContext::default())).await; + match expected { + ExpectedSubqueryResult::Value(expected) => { + assert_eq!(actual?, expected, "{name}"); + } + ExpectedSubqueryResult::Error(expected) => { + let err = actual.expect_err(name); + assert!( + err.to_string().contains(expected), + "{name}: expected error containing '{expected}', got {err}" + ); + } + } + } + + Ok(()) + } + + #[tokio::test] + async fn test_failed_subquery_is_not_retried() -> Result<()> { + let execute_calls = Arc::new(AtomicUsize::new(0)); + let subquery_plan = Arc::new(CountingExec::new( + Arc::new(ErrorExec::new()), + Arc::clone(&execute_calls), + )); + let exec = single_subquery_exec( + placeholder_input(), + subquery_plan, + ScalarSubqueryResults::new(1), + ); + + let ctx = Arc::new(TaskContext::default()); + let stream = exec.execute(0, Arc::clone(&ctx))?; + assert!(crate::common::collect(stream).await.is_err()); + + let stream = exec.execute(0, ctx)?; + assert!(crate::common::collect(stream).await.is_err()); + + assert_eq!(execute_calls.load(Ordering::SeqCst), 1); + Ok(()) + } + + #[tokio::test] + async fn test_reset_state_clears_results_and_reexecutes_subqueries() -> Result<()> { + let execute_calls = Arc::new(AtomicUsize::new(0)); + let results = ScalarSubqueryResults::new(1); + let subquery_plan = Arc::new(CountingExec::new( + make_subquery_plan(vec![int32_batch(vec![42])]), + Arc::clone(&execute_calls), + )); + let exec: Arc<dyn ExecutionPlan> = Arc::new(single_subquery_exec( + scalar_subquery_projection_input(results.clone())?, + subquery_plan, + results.clone(), + )); + + let batches = + crate::common::collect(exec.execute(0, Arc::new(TaskContext::default()))?) + .await?; + assert_eq!(extract_single_int32_value(&batches), 42); + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int32(Some(42))) + ); + + let reset_exec = reset_plan_states(Arc::clone(&exec))?; + assert_eq!(results.get(SubqueryIndex::new(0)), None); + + let reset_batches = crate::common::collect( + reset_exec.execute(0, Arc::new(TaskContext::default()))?, + ) + .await?; + assert_eq!(extract_single_int32_value(&reset_batches), 42); + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int32(Some(42))) + ); + assert_eq!(execute_calls.load(Ordering::SeqCst), 2); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sort_pushdown.rs b/native/vendor/datafusion-physical-plan/src/sort_pushdown.rs new file mode 100644 index 0000000..8432fd5 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sort_pushdown.rs
@@ -0,0 +1,120 @@ +// 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 pushdown types for physical execution plans. +//! +//! This module provides types used for pushing sort ordering requirements +//! down through the execution plan tree to data sources. + +/// Result of attempting to push down sort ordering to a node. +/// +/// Used by [`ExecutionPlan::try_pushdown_sort`] to communicate +/// whether and how sort ordering was successfully pushed down. +/// +/// [`ExecutionPlan::try_pushdown_sort`]: crate::ExecutionPlan::try_pushdown_sort +#[derive(Debug, Clone)] +pub enum SortOrderPushdownResult<T> { + /// The source can guarantee exact ordering (data is perfectly sorted). + /// + /// When this is returned, the optimizer can safely remove the Sort operator + /// entirely since the data source guarantees the requested ordering. + Exact { + /// The optimized node that provides exact ordering + inner: T, + }, + /// The source has optimized for the ordering but cannot guarantee perfect sorting. + /// + /// This indicates the data source has been optimized (e.g., reordered files/row groups + /// based on statistics, enabled reverse scanning) but the data may not be perfectly + /// sorted. The optimizer should keep the Sort operator but benefits from the + /// optimization (e.g., faster TopK queries due to early termination). + Inexact { + /// The optimized node that provides approximate ordering + inner: T, + }, + /// The source cannot optimize for this ordering. + /// + /// The data source does not support the requested sort ordering and no + /// optimization was applied. + Unsupported, +} + +impl<T> SortOrderPushdownResult<T> { + /// Extract the inner value if present + pub fn into_inner(self) -> Option<T> { + match self { + Self::Exact { inner } | Self::Inexact { inner } => Some(inner), + Self::Unsupported => None, + } + } + + /// Map the inner value to a different type while preserving the variant. + pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> SortOrderPushdownResult<U> { + match self { + Self::Exact { inner } => SortOrderPushdownResult::Exact { inner: f(inner) }, + Self::Inexact { inner } => { + SortOrderPushdownResult::Inexact { inner: f(inner) } + } + Self::Unsupported => SortOrderPushdownResult::Unsupported, + } + } + + /// Try to map the inner value, returning an error if the function fails. + pub fn try_map<U, E, F: FnOnce(T) -> Result<U, E>>( + self, + f: F, + ) -> Result<SortOrderPushdownResult<U>, E> { + match self { + Self::Exact { inner } => { + Ok(SortOrderPushdownResult::Exact { inner: f(inner)? }) + } + Self::Inexact { inner } => { + Ok(SortOrderPushdownResult::Inexact { inner: f(inner)? }) + } + Self::Unsupported => Ok(SortOrderPushdownResult::Unsupported), + } + } + + /// Convert this result to `Inexact`, downgrading `Exact` if present. + /// + /// This is useful when an operation (like merging multiple partitions) + /// cannot guarantee exact ordering even if the input provides it. + /// + /// # Examples + /// + /// ``` + /// # use datafusion_physical_plan::SortOrderPushdownResult; + /// let exact = SortOrderPushdownResult::Exact { inner: 42 }; + /// let inexact = exact.into_inexact(); + /// assert!(matches!(inexact, SortOrderPushdownResult::Inexact { inner: 42 })); + /// + /// let already_inexact = SortOrderPushdownResult::Inexact { inner: 42 }; + /// let still_inexact = already_inexact.into_inexact(); + /// assert!(matches!(still_inexact, SortOrderPushdownResult::Inexact { inner: 42 })); + /// + /// let unsupported = SortOrderPushdownResult::<i32>::Unsupported; + /// let still_unsupported = unsupported.into_inexact(); + /// assert!(matches!(still_unsupported, SortOrderPushdownResult::Unsupported)); + /// ``` + pub fn into_inexact(self) -> Self { + match self { + Self::Exact { inner } => Self::Inexact { inner }, + Self::Inexact { inner } => Self::Inexact { inner }, + Self::Unsupported => Self::Unsupported, + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/builder.rs b/native/vendor/datafusion-physical-plan/src/sorts/builder.rs new file mode 100644 index 0000000..75eb2ff --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/builder.rs
@@ -0,0 +1,359 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::spill::get_record_batch_memory_size; +use arrow::array::ArrayRef; +use arrow::compute::interleave; +use arrow::datatypes::SchemaRef; +use arrow::error::ArrowError; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::MemoryReservation; +use log::warn; +use std::sync::Arc; + +#[derive(Debug, Copy, Clone, Default)] +struct BatchCursor { + /// The index into BatchBuilder::batches + batch_idx: usize, + /// The row index within the given batch + row_idx: usize, +} + +/// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`] +#[derive(Debug)] +pub struct BatchBuilder { + /// The schema of the RecordBatches yielded by this stream + schema: SchemaRef, + + /// Maintain a list of [`RecordBatch`] and their corresponding stream + batches: Vec<(usize, RecordBatch)>, + + /// Accounts for memory used by buffered batches. + /// + /// May include pre-reserved bytes (from `sort_spill_reservation_bytes`) + /// that were transferred via [`MemoryReservation::take()`] to prevent + /// starvation when concurrent sort partitions compete for pool memory. + reservation: MemoryReservation, + + /// Tracks the actual memory used by buffered batches (not including + /// pre-reserved bytes). This allows [`Self::push_batch`] to skip pool + /// allocation requests when the pre-reserved bytes cover the batch. + batches_mem_used: usize, + + /// The initial reservation size at construction time. When the reservation + /// is pre-loaded with `sort_spill_reservation_bytes` (via `take()`), this + /// records that amount so we never shrink below it, maintaining the + /// anti-starvation guarantee throughout the merge. + initial_reservation: usize, + + /// The current [`BatchCursor`] for each stream + cursors: Vec<BatchCursor>, + + /// The accumulated stream indexes from which to pull rows + /// Consists of a tuple of `(batch_idx, row_idx)` + indices: Vec<(usize, usize)>, +} + +impl BatchBuilder { + /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size` + pub fn new( + schema: SchemaRef, + stream_count: usize, + batch_size: usize, + reservation: MemoryReservation, + ) -> Self { + let initial_reservation = reservation.size(); + Self { + schema, + batches: Vec::with_capacity(stream_count * 2), + cursors: vec![BatchCursor::default(); stream_count], + indices: Vec::with_capacity(batch_size), + reservation, + batches_mem_used: 0, + initial_reservation, + } + } + + /// Append a new batch in `stream_idx` + pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<()> { + let size = get_record_batch_memory_size(&batch); + self.batches_mem_used += size; + // Only request additional memory from the pool when actual batch + // usage exceeds the current reservation (which may include + // pre-reserved bytes from sort_spill_reservation_bytes). + try_grow_reservation_to_at_least(&mut self.reservation, self.batches_mem_used)?; + let batch_idx = self.batches.len(); + self.batches.push((stream_idx, batch)); + self.cursors[stream_idx] = BatchCursor { + batch_idx, + row_idx: 0, + }; + Ok(()) + } + + /// Append the next row from `stream_idx` + pub fn push_row(&mut self, stream_idx: usize) { + let cursor = &mut self.cursors[stream_idx]; + let row_idx = cursor.row_idx; + cursor.row_idx += 1; + self.indices.push((cursor.batch_idx, row_idx)); + } + + /// Returns the number of in-progress rows in this [`BatchBuilder`] + pub fn len(&self) -> usize { + self.indices.len() + } + + /// Returns `true` if this [`BatchBuilder`] contains no in-progress rows + pub fn is_empty(&self) -> bool { + self.indices.is_empty() + } + + /// Returns the schema of this [`BatchBuilder`] + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Try to interleave all columns using the given index slice. + fn try_interleave_columns( + &self, + indices: &[(usize, usize)], + ) -> Result<Vec<ArrayRef>> { + (0..self.schema.fields.len()) + .map(|column_idx| { + let arrays: Vec<_> = self + .batches + .iter() + .map(|(_, batch)| batch.column(column_idx).as_ref()) + .collect(); + // Arrow 58.1.0+ returns OffsetOverflowError directly from + // interleave, allowing retry_interleave to shrink the batch. + interleave(&arrays, indices).map_err(Into::into) + }) + .collect::<Result<Vec<_>>>() + } + + /// Builds a record batch from the first `rows_to_emit` buffered rows. + fn finish_record_batch( + &mut self, + rows_to_emit: usize, + columns: Vec<ArrayRef>, + ) -> Result<RecordBatch> { + // Remove consumed indices, keeping any remaining for the next call. + self.indices.drain(..rows_to_emit); + + // Only clean up fully-consumed batches when all indices are drained, + // because remaining indices may still reference earlier batches. + // In the overflow/partial-emit case this may retain some extra memory + // across a few drain polls, but avoids costly index scanning on the + // hot path. The retention is bounded and short-lived since leftover + // rows are drained over subsequent polls. + if self.indices.is_empty() { + // New cursors are only created once the previous cursor for the stream + // is finished. This means all remaining rows from all but the last batch + // for each stream have been yielded to the newly created record batch + // + // We can therefore drop all but the last batch for each stream + let mut batch_idx = 0; + let mut retained = 0; + self.batches.retain(|(stream_idx, batch)| { + let stream_cursor = &mut self.cursors[*stream_idx]; + let retain = stream_cursor.batch_idx == batch_idx; + batch_idx += 1; + + if retain { + stream_cursor.batch_idx = retained; + retained += 1; + } else { + self.batches_mem_used -= get_record_batch_memory_size(batch); + } + retain + }); + } + + // Release excess memory back to the pool, but never shrink below + // initial_reservation to maintain the anti-starvation guarantee + // for the merge phase. + let target = self.batches_mem_used.max(self.initial_reservation); + if self.reservation.size() > target { + self.reservation.shrink(self.reservation.size() - target); + } + + RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) + } + + /// Drains the in_progress row indexes, and builds a new RecordBatch from them + /// + /// Will then drop any batches for which all rows have been yielded to the output. + /// If an offset overflow occurs (e.g. string/list offsets exceed i32::MAX), + /// retries with progressively fewer rows until it succeeds. + /// + /// Returns `None` if no pending rows + pub fn build_record_batch(&mut self) -> Result<Option<RecordBatch>> { + if self.is_empty() { + return Ok(None); + } + + let (rows_to_emit, columns) = + retry_interleave(self.indices.len(), self.indices.len(), |rows_to_emit| { + self.try_interleave_columns(&self.indices[..rows_to_emit]) + })?; + + Ok(Some(self.finish_record_batch(rows_to_emit, columns)?)) + } +} + +/// Try to grow `reservation` so it covers at least `needed` bytes. +/// +/// When a reservation has been pre-loaded with bytes (e.g. via +/// [`MemoryReservation::take()`]), this avoids redundant pool +/// allocations: if the reservation already covers `needed`, this is +/// a no-op; otherwise only the deficit is requested from the pool. +pub(crate) fn try_grow_reservation_to_at_least( + reservation: &mut MemoryReservation, + needed: usize, +) -> Result<()> { + if needed > reservation.size() { + reservation.try_grow(needed - reservation.size())?; + } + Ok(()) +} + +/// Returns true if the error is an Arrow offset overflow. +fn is_offset_overflow(e: &DataFusionError) -> bool { + matches!( + e, + DataFusionError::ArrowError(boxed, _) + if matches!(boxed.as_ref(), ArrowError::OffsetOverflowError(_)) + ) +} + +#[cfg(test)] +fn offset_overflow_error() -> DataFusionError { + DataFusionError::ArrowError(Box::new(ArrowError::OffsetOverflowError(0)), None) +} + +fn retry_interleave<T, F>( + mut rows_to_emit: usize, + total_rows: usize, + mut interleave: F, +) -> Result<(usize, T)> +where + F: FnMut(usize) -> Result<T>, +{ + loop { + match interleave(rows_to_emit) { + Ok(value) => return Ok((rows_to_emit, value)), + // Only offset overflow is recoverable by emitting fewer rows. + Err(e) if is_offset_overflow(&e) => { + rows_to_emit /= 2; + if rows_to_emit == 0 { + return Err(e); + } + warn!( + "Interleave offset overflow with {total_rows} rows, retrying with {rows_to_emit}" + ); + } + Err(e) => return Err(e), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, ArrayDataBuilder, Int32Array, ListArray}; + use arrow::buffer::Buffer; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryPool, UnboundedMemoryPool, + }; + + fn overflow_list_batch() -> RecordBatch { + let values_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + // SAFETY: This intentionally constructs an invalid child length so + // Arrow's interleave hits offset overflow before touching child data. + let list = ListArray::from(unsafe { + ArrayDataBuilder::new(DataType::List(Arc::clone(&values_field))) + .len(1) + .add_buffer(Buffer::from_slice_ref([0_i32, i32::MAX])) + .add_child_data(Int32Array::from(Vec::<i32>::new()).to_data()) + .build_unchecked() + }); + let schema = Arc::new(Schema::new(vec![Field::new( + "list_col", + DataType::List(values_field), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(list)]).unwrap() + } + + #[test] + fn test_retry_interleave_halves_rows_until_success() { + let mut attempts = Vec::new(); + + let (rows_to_emit, result) = retry_interleave(4, 4, |rows_to_emit| { + attempts.push(rows_to_emit); + if rows_to_emit > 1 { + Err(offset_overflow_error()) + } else { + Ok("ok") + } + }) + .unwrap(); + + assert_eq!(rows_to_emit, 1); + assert_eq!(result, "ok"); + assert_eq!(attempts, vec![4, 2, 1]); + } + + #[test] + fn test_is_offset_overflow_matches_arrow_error() { + assert!(is_offset_overflow(&offset_overflow_error())); + } + + #[test] + fn test_retry_interleave_does_not_retry_non_offset_errors() { + let mut attempts = Vec::new(); + + let error = retry_interleave(4, 4, |rows_to_emit| { + attempts.push(rows_to_emit); + Err::<(), _>(DataFusionError::Execution("boom".into())) + }) + .unwrap_err(); + + assert_eq!(attempts, vec![4]); + assert!(matches!(error, DataFusionError::Execution(msg) if msg == "boom")); + } + + #[test] + fn test_try_interleave_columns_surfaces_arrow_offset_overflow() { + let batch = overflow_list_batch(); + let schema = batch.schema(); + let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("test").register(&pool); + let mut builder = BatchBuilder::new(schema, 1, 2, reservation); + builder.push_batch(0, batch).unwrap(); + + let error = builder + .try_interleave_columns(&[(0, 0), (0, 0)]) + .unwrap_err(); + + assert!(is_offset_overflow(&error)); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/cursor.rs b/native/vendor/datafusion-physical-plan/src/sorts/cursor.rs new file mode 100644 index 0000000..d71eaad --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/cursor.rs
@@ -0,0 +1,691 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cmp::Ordering; +use std::fmt::Debug; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrowPrimitiveType, GenericByteArray, GenericByteViewArray, OffsetSizeTrait, + PrimitiveArray, StringViewArray, types::ByteArrayType, +}; +use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; +use arrow::compute::SortOptions; +use arrow::datatypes::ArrowNativeTypeOp; +use arrow::row::Rows; +use datafusion_execution::memory_pool::MemoryReservation; + +/// A comparable collection of values for use with [`Cursor`] +/// +/// This is a trait as there are several specialized implementations, such as for +/// single columns or for normalized multi column keys ([`Rows`]) +pub trait CursorValues: Debug + Sync + Send { + fn len(&self) -> usize; + + /// Returns true if `l[l_idx] == r[r_idx]` + fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool; + + /// Returns true if `row[idx] == row[idx - 1]` + /// Given `idx` should be greater than 0 + fn eq_to_previous(cursor: &Self, idx: usize) -> bool; + + /// Returns comparison of `l[l_idx]` and `r[r_idx]` + fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; + + /// Notifies the values that the owning [`Cursor`] moved to `offset` (always + /// `< len()`), so caching implementations can refresh the value(s) read by + /// the hot comparisons. Default no-op (e.g. byte/row cursors don't benefit). + #[inline] + fn set_offset(&mut self, offset: usize) { + let _ = offset; + } +} + +/// A comparable cursor, used by sort operations +/// +/// A `Cursor` is a pointer into a collection of rows, stored in +/// [`CursorValues`] +/// +/// ```text +/// +/// ┌───────────────────────┐ +/// │ │ ┌──────────────────────┐ +/// │ ┌─────────┐ ┌─────┐ │ ─ ─ ─ ─│ Cursor<T> │ +/// │ │ 1 │ │ A │ │ │ └──────────────────────┘ +/// │ ├─────────┤ ├─────┤ │ +/// │ │ 2 │ │ A │◀─ ┼ ─ ┘ Cursor<T> tracks an +/// │ └─────────┘ └─────┘ │ offset within a +/// │ ... ... │ CursorValues +/// │ │ +/// │ ┌─────────┐ ┌─────┐ │ +/// │ │ 3 │ │ E │ │ +/// │ └─────────┘ └─────┘ │ +/// │ │ +/// │ CursorValues │ +/// └───────────────────────┘ +/// ``` +/// +/// Store logical rows using one of several formats, with specialized +/// implementations depending on the column types +#[derive(Debug)] +pub struct Cursor<T: CursorValues> { + offset: usize, + values: T, +} + +impl<T: CursorValues> Cursor<T> { + /// Create a [`Cursor`] from the given [`CursorValues`] + pub fn new(values: T) -> Self { + Self { offset: 0, values } + } + + /// Returns true if there are no more rows in this cursor + #[inline] + pub fn is_finished(&self) -> bool { + self.offset == self.values.len() + } + + /// Advance the cursor, returning the previous row index + #[inline] + pub fn advance(&mut self) -> usize { + let t = self.offset; + self.offset += 1; + // Refresh the cache for the new position. The guard keeps `set_offset` + // in bounds; a finished cursor's stale cache is never read (it is taken + // before the next comparison). + if self.offset < self.values.len() { + self.values.set_offset(self.offset); + } + t + } + + pub fn is_eq_to_prev_one(&self, prev_cursor: Option<&Cursor<T>>) -> bool { + if self.offset > 0 { + self.is_eq_to_prev_row() + } else if let Some(prev_cursor) = prev_cursor { + self.is_eq_to_prev_row_in_prev_batch(prev_cursor) + } else { + false + } + } +} + +impl<T: CursorValues> PartialEq for Cursor<T> { + #[inline] + fn eq(&self, other: &Self) -> bool { + T::eq(&self.values, self.offset, &other.values, other.offset) + } +} + +impl<T: CursorValues> Cursor<T> { + fn is_eq_to_prev_row(&self) -> bool { + T::eq_to_previous(&self.values, self.offset) + } + + fn is_eq_to_prev_row_in_prev_batch(&self, other: &Self) -> bool { + assert_eq!(self.offset, 0); + T::eq( + &self.values, + self.offset, + &other.values, + other.values.len() - 1, + ) + } +} + +impl<T: CursorValues> Eq for Cursor<T> {} + +impl<T: CursorValues> PartialOrd for Cursor<T> { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + Some(self.cmp(other)) + } +} + +impl<T: CursorValues> Ord for Cursor<T> { + #[inline] + fn cmp(&self, other: &Self) -> Ordering { + T::compare(&self.values, self.offset, &other.values, other.offset) + } +} + +/// Implements [`CursorValues`] for [`Rows`] +/// +/// Used for sorting when there are multiple columns in the sort key +#[derive(Debug)] +pub struct RowValues { + rows: Arc<Rows>, + + /// Tracks for the memory used by in the `Rows` of this + /// cursor. Freed on drop + _reservation: MemoryReservation, +} + +impl RowValues { + /// Create a new [`RowValues`] from `rows` and a `reservation` + /// that tracks its memory. There must be at least one row + /// + /// Panics if the reservation is not for exactly `rows.size()` + /// bytes or if `rows` is empty. + pub fn new(rows: Arc<Rows>, reservation: MemoryReservation) -> Self { + assert_eq!( + rows.size(), + reservation.size(), + "memory reservation mismatch" + ); + assert!(rows.num_rows() > 0); + Self { + rows, + _reservation: reservation, + } + } +} + +impl CursorValues for RowValues { + #[inline] + fn len(&self) -> usize { + self.rows.num_rows() + } + + // No inline hint on purpose: for the heavyweight `Rows` byte comparison the + // compiler's own choice wins — both `#[inline]` and `#[inline(never)]` + // measurably regress the multi-column merge path. + fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { + l.rows.row(l_idx) == r.rows.row(r_idx) + } + + fn eq_to_previous(cursor: &Self, idx: usize) -> bool { + assert!(idx > 0); + cursor.rows.row(idx) == cursor.rows.row(idx - 1) + } + + fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + l.rows.row(l_idx).cmp(&r.rows.row(r_idx)) + } +} + +/// An [`Array`] that can be converted into [`CursorValues`] +pub trait CursorArray: Array + 'static { + type Values: CursorValues; + + fn values(&self) -> Self::Values; +} + +impl<T: ArrowPrimitiveType> CursorArray for PrimitiveArray<T> { + type Values = PrimitiveValues<T::Native>; + + fn values(&self) -> Self::Values { + PrimitiveValues::new(self.values().clone()) + } +} + +/// [`CursorValues`] for a primitive column. +/// +/// Caches the value at the current (and previous) offset, refreshed once per +/// [`Cursor::advance`] via [`CursorValues::set_offset`], so the hot loser-tree +/// comparisons read a cached field instead of indexing the buffer each time. +#[derive(Debug)] +pub struct PrimitiveValues<T: ArrowNativeTypeOp> { + values: ScalarBuffer<T>, + /// Cached `values[offset]`. + current: T, + /// Cached `values[offset - 1]` (read by `eq_to_previous`, only past offset 0). + previous: T, + /// Current offset; used only to `debug_assert!` the cache is read in sync. + offset: usize, +} + +impl<T: ArrowNativeTypeOp> PrimitiveValues<T> { + fn new(values: ScalarBuffer<T>) -> Self { + // Non-empty in practice; `unwrap_or_default` just avoids a panic. + let first = values.first().copied().unwrap_or_default(); + Self { + values, + current: first, + previous: first, + offset: 0, + } + } +} + +impl<T: ArrowNativeTypeOp> CursorValues for PrimitiveValues<T> { + #[inline(always)] + fn len(&self) -> usize { + self.values.len() + } + + #[inline(always)] + fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { + // Arbitrary indices (cross-batch comparison), so index directly. + l.values[l_idx].is_eq(r.values[r_idx]) + } + + #[inline(always)] + fn eq_to_previous(cursor: &Self, idx: usize) -> bool { + assert!(idx > 0); + debug_assert_eq!(idx, cursor.offset); + cursor.current.is_eq(cursor.previous) + } + + #[inline(always)] + fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + debug_assert_eq!(l_idx, l.offset); + debug_assert_eq!(r_idx, r.offset); + l.current.compare(r.current) + } + + #[inline(always)] + fn set_offset(&mut self, offset: usize) { + // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that + // guard dominates the index below so its bounds check is elided — the + // length is checked once per row, not per comparison. The old `current` + // is `values[offset - 1]`, so it becomes `previous`. + self.previous = self.current; + self.current = self.values[offset]; + self.offset = offset; + } +} + +#[derive(Debug)] +pub struct ByteArrayValues<T: OffsetSizeTrait> { + offsets: OffsetBuffer<T>, + values: Buffer, +} + +impl<T: OffsetSizeTrait> ByteArrayValues<T> { + #[inline] + fn value(&self, idx: usize) -> &[u8] { + assert!(idx < self.len()); + // Safety: offsets are valid and checked bounds above + unsafe { + let start = self.offsets.get_unchecked(idx).as_usize(); + let end = self.offsets.get_unchecked(idx + 1).as_usize(); + self.values.get_unchecked(start..end) + } + } +} + +impl<T: OffsetSizeTrait> CursorValues for ByteArrayValues<T> { + #[inline] + fn len(&self) -> usize { + self.offsets.len() - 1 + } + + #[inline] + fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { + l.value(l_idx) == r.value(r_idx) + } + + #[inline] + fn eq_to_previous(cursor: &Self, idx: usize) -> bool { + assert!(idx > 0); + cursor.value(idx) == cursor.value(idx - 1) + } + + #[inline] + fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + l.value(l_idx).cmp(r.value(r_idx)) + } +} + +impl<T: ByteArrayType> CursorArray for GenericByteArray<T> { + type Values = ByteArrayValues<T::Offset>; + + fn values(&self) -> Self::Values { + ByteArrayValues { + offsets: self.offsets().clone(), + values: self.values().clone(), + } + } +} + +impl CursorArray for StringViewArray { + type Values = StringViewArray; + fn values(&self) -> Self { + self.gc() + } +} + +impl CursorValues for StringViewArray { + fn len(&self) -> usize { + self.views().len() + } + + #[inline(always)] + fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { + // SAFETY: Both l_idx and r_idx are guaranteed to be within bounds, + // and any null-checks are handled in the outer layers. + // Fast path: Compare the lengths before full byte comparison. + let l_view = unsafe { l.views().get_unchecked(l_idx) }; + let r_view = unsafe { r.views().get_unchecked(r_idx) }; + + if l.data_buffers().is_empty() && r.data_buffers().is_empty() { + return l_view == r_view; + } + + let l_len = *l_view as u32; + let r_len = *r_view as u32; + if l_len != r_len { + return false; + } + + unsafe { GenericByteViewArray::compare_unchecked(l, l_idx, r, r_idx).is_eq() } + } + + #[inline(always)] + fn eq_to_previous(cursor: &Self, idx: usize) -> bool { + // SAFETY: The caller guarantees that idx > 0 and the indices are valid. + // Already checked it in is_eq_to_prev_one function + // Fast path: Compare the lengths of the current and previous views. + let l_view = unsafe { cursor.views().get_unchecked(idx) }; + let r_view = unsafe { cursor.views().get_unchecked(idx - 1) }; + if cursor.data_buffers().is_empty() { + return l_view == r_view; + } + + let l_len = *l_view as u32; + let r_len = *r_view as u32; + + if l_len != r_len { + return false; + } + + unsafe { + GenericByteViewArray::compare_unchecked(cursor, idx, cursor, idx - 1).is_eq() + } + } + + #[inline(always)] + fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + // SAFETY: Prior assertions guarantee that l_idx and r_idx are valid indices. + // Null-checks are assumed to have been handled in the wrapper (e.g., ArrayValues). + // And the bound is checked in is_finished, it is safe to call get_unchecked + if l.data_buffers().is_empty() && r.data_buffers().is_empty() { + let l_view = unsafe { l.views().get_unchecked(l_idx) }; + let r_view = unsafe { r.views().get_unchecked(r_idx) }; + return StringViewArray::inline_key_fast(*l_view) + .cmp(&StringViewArray::inline_key_fast(*r_view)); + } + + unsafe { GenericByteViewArray::compare_unchecked(l, l_idx, r, r_idx) } + } +} + +/// A collection of sorted, nullable [`CursorValues`] +/// +/// Note: comparing cursors with different `SortOptions` will yield an arbitrary ordering +#[derive(Debug)] +pub struct ArrayValues<T: CursorValues> { + values: T, + // If nulls first, the first non-null index + // Otherwise, the first null index + null_threshold: usize, + options: SortOptions, + + /// Tracks the memory used by the values array, + /// freed on drop. + _reservation: MemoryReservation, +} + +impl<T: CursorValues> ArrayValues<T> { + /// Create a new [`ArrayValues`] from the provided `values` sorted according + /// to `options`. + /// + /// Panics if the array is empty + pub fn new<A: CursorArray<Values = T>>( + options: SortOptions, + array: &A, + reservation: MemoryReservation, + ) -> Self { + assert!(array.len() > 0, "Empty array passed to FieldCursor"); + let null_threshold = match options.nulls_first { + true => array.null_count(), + false => array.len() - array.null_count(), + }; + + Self { + values: array.values(), + null_threshold, + options, + _reservation: reservation, + } + } + + #[inline(always)] + fn is_null(&self, idx: usize) -> bool { + (idx < self.null_threshold) == self.options.nulls_first + } +} + +impl<T: CursorValues> CursorValues for ArrayValues<T> { + #[inline(always)] + fn len(&self) -> usize { + self.values.len() + } + + #[inline(always)] + fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { + match (l.is_null(l_idx), r.is_null(r_idx)) { + (true, true) => true, + (false, false) => T::eq(&l.values, l_idx, &r.values, r_idx), + _ => false, + } + } + + #[inline(always)] + fn eq_to_previous(cursor: &Self, idx: usize) -> bool { + assert!(idx > 0); + match (cursor.is_null(idx), cursor.is_null(idx - 1)) { + (true, true) => true, + // Delegate to inner `eq_to_previous` so a caching cursor can answer + // without indexing. + (false, false) => T::eq_to_previous(&cursor.values, idx), + _ => false, + } + } + + #[inline(always)] + fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { + match (l.is_null(l_idx), r.is_null(r_idx)) { + (true, true) => Ordering::Equal, + (true, false) => match l.options.nulls_first { + true => Ordering::Less, + false => Ordering::Greater, + }, + (false, true) => match l.options.nulls_first { + true => Ordering::Greater, + false => Ordering::Less, + }, + (false, false) => match l.options.descending { + true => T::compare(&r.values, r_idx, &l.values, l_idx), + false => T::compare(&l.values, l_idx, &r.values, r_idx), + }, + } + } + + #[inline(always)] + fn set_offset(&mut self, offset: usize) { + // Forward to the wrapped values (e.g. caching `PrimitiveValues`). + self.values.set_offset(offset); + } +} + +#[cfg(test)] +mod tests { + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; + + use super::*; + + fn new_primitive( + options: SortOptions, + values: ScalarBuffer<i32>, + null_count: usize, + ) -> Cursor<ArrayValues<PrimitiveValues<i32>>> { + let null_threshold = match options.nulls_first { + true => null_count, + false => values.len() - null_count, + }; + + let memory_pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(10000)); + let consumer = MemoryConsumer::new("test"); + let reservation = consumer.register(&memory_pool); + + let values = ArrayValues { + values: PrimitiveValues::new(values), + null_threshold, + options, + _reservation: reservation, + }; + + Cursor::new(values) + } + + #[test] + fn test_primitive_nulls_first() { + let options = SortOptions { + descending: false, + nulls_first: true, + }; + + let buffer = ScalarBuffer::from(vec![i32::MAX, 1, 2, 3]); + let mut a = new_primitive(options, buffer, 1); + let buffer = ScalarBuffer::from(vec![1, 2, -2, -1, 1, 9]); + let mut b = new_primitive(options, buffer, 2); + + // NULL == NULL + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + // NULL == NULL + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + // NULL < -2 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // 1 > -2 + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Greater); + + // 1 > -1 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Greater); + + // 1 == 1 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + // 9 > 1 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // 9 > 2 + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + let options = SortOptions { + descending: false, + nulls_first: false, + }; + + let buffer = ScalarBuffer::from(vec![0, 1, i32::MIN, i32::MAX]); + let mut a = new_primitive(options, buffer, 2); + let buffer = ScalarBuffer::from(vec![-1, i32::MAX, i32::MIN]); + let mut b = new_primitive(options, buffer, 2); + + // 0 > -1 + assert_eq!(a.cmp(&b), Ordering::Greater); + + // 0 < NULL + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // 1 < NULL + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // NULL = NULL + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + let options = SortOptions { + descending: true, + nulls_first: false, + }; + + let buffer = ScalarBuffer::from(vec![6, 1, i32::MIN, i32::MAX]); + let mut a = new_primitive(options, buffer, 3); + let buffer = ScalarBuffer::from(vec![67, -3, i32::MAX, i32::MIN]); + let mut b = new_primitive(options, buffer, 2); + + // 6 > 67 + assert_eq!(a.cmp(&b), Ordering::Greater); + + // 6 < -3 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // 6 < NULL + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // 6 < NULL + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // NULL == NULL + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + let options = SortOptions { + descending: true, + nulls_first: true, + }; + + let buffer = ScalarBuffer::from(vec![i32::MIN, i32::MAX, 6, 3]); + let mut a = new_primitive(options, buffer, 2); + let buffer = ScalarBuffer::from(vec![i32::MAX, 4546, -3]); + let mut b = new_primitive(options, buffer, 1); + + // NULL == NULL + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + // NULL == NULL + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Equal); + assert_eq!(a, b); + + // NULL < 4546 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + + // 6 > 4546 + a.advance(); + assert_eq!(a.cmp(&b), Ordering::Greater); + + // 6 < -3 + b.advance(); + assert_eq!(a.cmp(&b), Ordering::Less); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/merge.rs b/native/vendor/datafusion-physical-plan/src/sorts/merge.rs new file mode 100644 index 0000000..6476490 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/merge.rs
@@ -0,0 +1,729 @@ +// 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. + +//! Merge that deals with an arbitrary size of streaming inputs. +//! This is an order-preserving merge. + +use std::fmt::Debug; +use std::future::poll_fn; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use crate::SendableRecordBatchStream; +use crate::metrics::BaselineMetrics; +use crate::sorts::builder::BatchBuilder; +use crate::sorts::cursor::{Cursor, CursorValues}; +use crate::sorts::stream::PartitionedStream; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result, assert_or_internal_err, internal_err}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{TryEmitter, async_try_stream}; +use futures::Stream; + +/// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] +type CursorStream<C> = Box<dyn PartitionedStream<Output = Result<(C, RecordBatch)>>>; + +/// Merges a stream of sorted cursors and record batches into a single sorted stream +#[derive(Debug)] +pub(crate) struct SortPreservingMergeStream<C: CursorValues> { + in_progress: BatchBuilder, + + /// The sorted input streams to merge together + streams: CursorStream<C>, + + /// used to record execution metrics + metrics: BaselineMetrics, + + /// A loser tree that always produces the minimum cursor + /// + /// Node 0 stores the top winner, Nodes 1..num_streams store + /// the loser nodes + /// + /// This implements a "Tournament Tree" (aka Loser Tree) to keep + /// track of the current smallest element at the top. When the top + /// record is taken, the tree structure is not modified, and only + /// the path from bottom to top is visited, keeping the number of + /// comparisons close to the theoretical limit of `log(S)`. + /// + /// The current implementation uses a vector to store the tree. + /// Conceptually, it looks like this (assuming 8 streams): + /// + /// ```text + /// 0 (winner) + /// + /// 1 + /// / \ + /// 2 3 + /// / \ / \ + /// 4 5 6 7 + /// ``` + /// + /// Where element at index 0 in the vector is the current winner. Element + /// at index 1 is the root of the loser tree, element at index 2 is the + /// left child of the root, and element at index 3 is the right child of + /// the root and so on. + /// + /// reference: <https://en.wikipedia.org/wiki/K-way_merge_algorithm#Tournament_Tree> + loser_tree: Vec<usize>, + + /// Target batch size + batch_size: usize, + + /// Cursors for each input partition. `None` means the input is exhausted + cursors: Vec<Option<Cursor<C>>>, + + /// Flag indicating whether we are in the mode of round-robin + /// tie breaker for the loser tree winners. + round_robin_tie_breaker_mode: bool, + + /// Total number of polls returning the same value, as per partition. + /// We select the one that has less poll counts for tie-breaker in loser tree. + num_of_polled_with_same_value: Vec<usize>, + + /// To keep track of reset counts + poll_reset_epochs: Vec<usize>, + + /// Current reset count + current_reset_epoch: usize, + + /// Stores the previous value of each partitions for tracking the poll counts on the same value + /// Used if and only if round robin tie breaker is enabled, otherwise None + prev_cursors: Option<Vec<Option<Cursor<C>>>>, + + /// Optional number of rows to fetch + fetch: Option<usize>, + + /// number of rows produced + produced: usize, +} + +impl<C: CursorValues> SortPreservingMergeStream<C> { + pub(crate) fn new( + streams: CursorStream<C>, + schema: SchemaRef, + metrics: BaselineMetrics, + batch_size: usize, + fetch: Option<usize>, + reservation: MemoryReservation, + enable_round_robin_tie_breaker: bool, + ) -> Self { + assert_ne!(batch_size, 0, "batch size cannot be 0"); + assert_ne!(fetch, Some(0), "fetch must not be Some(0)"); + + let stream_count = streams.partitions(); + + Self { + in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation), + streams, + metrics, + cursors: (0..stream_count).map(|_| None).collect(), + prev_cursors: if enable_round_robin_tie_breaker { + Some((0..stream_count).map(|_| None).collect()) + } else { + None + }, + round_robin_tie_breaker_mode: false, + num_of_polled_with_same_value: vec![0; stream_count], + current_reset_epoch: 0, + poll_reset_epochs: vec![0; stream_count], + loser_tree: vec![], + batch_size, + fetch, + produced: 0, + } + } + + pub(crate) fn into_stream(self) -> SendableRecordBatchStream + where + C: 'static, + { + let schema_clone = Arc::clone(self.in_progress.schema()); + + let cloned_metrics = self.metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + + /// If the stream at the given index is not exhausted, and the last cursor for the + /// stream is finished, poll the stream for the next RecordBatch and create a new + /// cursor for the stream from the returned result + fn maybe_poll_stream( + &mut self, + cx: &mut Context<'_>, + idx: usize, + ) -> Poll<Result<()>> { + if self.cursors[idx].is_some() { + // Cursor is not finished - don't need a new RecordBatch yet + return Poll::Ready(Ok(())); + } + + match futures::ready!(self.streams.poll_next(cx, idx)) { + None => Poll::Ready(Ok(())), + Some(Err(e)) => Poll::Ready(Err(e)), + Some(Ok((cursor, batch))) => { + self.cursors[idx] = Some(Cursor::new(cursor)); + Poll::Ready(self.in_progress.push_batch(idx, batch)) + } + } + } + + fn emit_in_progress_batch(&mut self) -> Result<Option<RecordBatch>> { + let rows_before = self.in_progress.len(); + let result = self.in_progress.build_record_batch(); + self.produced += rows_before - self.in_progress.len(); + result + } + + async fn flush_in_progress( + &mut self, + mut emitter: TryEmitter<RecordBatch, DataFusionError>, + ) -> Result<()> { + if self.in_progress.is_empty() { + return Ok(()); + } + + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + // When `build_record_batch()` hits an i32 offset overflow (e.g. + // combined string offsets exceed 2 GB), it emits a partial batch + // and keeps the remaining rows in `self.in_progress.indices`. + // Drain those leftover rows before terminating the stream, + // otherwise they would be silently dropped. + // Repeated overflows are fine — each poll emits another partial + // batch until `in_progress` is fully drained. + while let Some(batch) = self.emit_in_progress_batch()? { + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + Ok(()) + } + + fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> { + async_try_stream(|mut emitter| async move { + // 1. Make sure we have data from each stream so we can initialize the loser tree + { + // This vector contains the indices of the partitions that have not started emitting yet. + let mut uninitiated_partitions = + (0..self.streams.partitions()).collect::<Vec<_>>(); + + poll_fn(|cx| { + self.initialize_all_partitions(&mut uninitiated_partitions, cx) + }) + .await?; + + assert_eq!(uninitiated_partitions.len(), 0); + } + + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + // 2. Init loser tree + self.init_loser_tree(); + + // 3. loop until all streams have been exhausted + while !self.is_exhausted() { + // 3.1. add loser_tree[0] (minimum) stream to pending record batch + let winner_stream = self.loser_tree[0]; + self.in_progress.push_row(winner_stream); + + // 3.2. If the new row reached the limit + if self.fetch_reached() { + break; + } + + // 3.3. if there is enough to emit for a full record batch + if self.in_progress.len() >= self.batch_size { + // 3.3.1 build pending record batch and reset builder + let Some(batch) = self.emit_in_progress_batch()? else { + return internal_err!("must have batch in progress to emit"); + }; + + // 3.3.2 emit pending record batch + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + // 3.4. advance cursor for the winner stream + { + let should_poll_next_batch_for_stream = + self.advance_cursors(winner_stream); + + // Fast path: skip the `maybe_poll_stream` call (and its `Poll` + // plumbing) unless the winner's cursor is exhausted and needs a + // fresh batch — it is live for almost every row. + if should_poll_next_batch_for_stream { + assert_or_internal_err!( + self.cursors[winner_stream].is_none(), + "cursor should be exhausted" + ); + + drop(timer); + poll_fn(|cx| self.maybe_poll_stream(cx, winner_stream)).await?; + timer = elapsed_compute.timer(); + } + } + + // 3.5. Adjusting the loser tree if necessary + self.update_loser_tree(); + } + + // 4. Flush any remaining rows in `self.in_progress` + self.flush_in_progress(emitter).await?; + + Ok(()) + }) + } + + /// Returns `true` once every input stream is exhausted. + /// + /// Should only be called for valid adjusted tree, i.e. the initial tree or after [`Self::update_loser_tree`] call + fn is_exhausted(&self) -> bool { + let winner = self.loser_tree[0]; + + // Checking only the tree root suffices for valid tree + // since the winner of the tree cannot be an exhausted stream for a valid tree + // as what value is winning over the non exhausted stream? + self.cursors[winner].is_none() + } + + /// Initialize all partitions, return `Poll::Pending` if any partition returns `Poll::Pending` + /// + /// This DOES NOT return `Poll::Pending` as soon as the first uninitiated partition returns `Poll::Pending` + /// so we can continue to initialize the remaining partitions + fn initialize_all_partitions( + &mut self, + uninitiated_partitions: &mut Vec<usize>, + cx: &mut Context, + ) -> Poll<Result<()>> { + assert_eq!( + self.loser_tree.len(), + 0, + "loser tree must be empty when initializing" + ); + + // Manual indexing since we're iterating over the vector and shrinking it in the loop + let mut idx = 0; + while idx < uninitiated_partitions.len() { + let partition_idx = uninitiated_partitions[idx]; + match self.maybe_poll_stream(cx, partition_idx) { + Poll::Ready(Err(e)) => { + return Poll::Ready(Err(e)); + } + Poll::Pending => { + // The polled stream is pending which means we're already set up to + // be woken when necessary + // Try the next stream + idx += 1; + } + _ => { + // The polled stream is ready + // Remove it from uninitiated_partitions + // Don't bump idx here, since a new element will have taken its + // place which we'll try in the next loop iteration + // swap_remove will change the partition poll order, but that shouldn't + // make a difference since we're waiting for all streams to be ready. + uninitiated_partitions.swap_remove(idx); + } + } + } + + if uninitiated_partitions.is_empty() { + Poll::Ready(Ok(())) + } else { + // There are still uninitiated partitions so return pending. + // We only get here if we've polled all uninitiated streams and at least one of them + // returned pending itself. That means we will be woken as soon as one of the + // streams would like to be polled again. + // There is no need to reschedule ourselves eagerly. + Poll::Pending + } + } + + /// For the given partition, updates the poll count. If the current value is the same + /// of the previous value, it increases the count by 1; otherwise, it is reset as 0. + fn update_poll_count_on_the_same_value(&mut self, partition_idx: usize) { + let cursor = &mut self.cursors[partition_idx]; + + // Check if the current partition's poll count is logically "reset" + if self.poll_reset_epochs[partition_idx] != self.current_reset_epoch { + self.poll_reset_epochs[partition_idx] = self.current_reset_epoch; + self.num_of_polled_with_same_value[partition_idx] = 0; + } + + if let Some(c) = cursor.as_mut() { + // Compare with the last row in the previous batch + let prev_cursor = self + .prev_cursors + .as_ref() + .map(|v| &v[partition_idx]) + .expect( + "prev_cursor should be set when round robin tie breaker is enabled", + ); + if c.is_eq_to_prev_one(prev_cursor.as_ref()) { + self.num_of_polled_with_same_value[partition_idx] += 1; + } else { + self.num_of_polled_with_same_value[partition_idx] = 0; + } + } + } + + /// Whether round-robin selection of tied winners of loser tree is enabled. + /// + /// This option controls the tie-breaker strategy and attempts to avoid the + /// issue of unbalanced polling between partitions + /// + /// If `true`, when multiple partitions have the same value, the partition + /// that has the fewest poll counts is selected. This strategy ensures that + /// multiple partitions with the same value are chosen equally, distributing + /// the polling load in a round-robin fashion. This approach balances the + /// workload more effectively across partitions and avoids excessive buffer + /// growth. + /// + /// if `false`, partitions with smaller indices are consistently chosen as + /// the winners, which can lead to an uneven distribution of polling and potentially + /// causing upstream operator buffers for the other partitions to grow + /// excessively, as they continued receiving data without consuming it. + /// + /// For example, an upstream operator like `RepartitionExec` execution would + /// keep sending data to certain partitions, but those partitions wouldn't + /// consume the data if they weren't selected as winners. This resulted in + /// inefficient buffer usage. + fn round_robin_tie_breaker_enabled(&self) -> bool { + self.prev_cursors.is_some() + } + + fn fetch_reached(&mut self) -> bool { + self.fetch + .map(|fetch| self.produced + self.in_progress.len() >= fetch) + .unwrap_or(false) + } + + /// Advances the actual cursor. If it reaches its end, update the + /// previous cursor with it. + /// + /// If the given partition batch is exhausted, return `true` to signal a poll is needed + fn advance_cursors(&mut self, stream_idx: usize) -> bool { + if let Some(cursor) = &mut self.cursors[stream_idx] { + let _ = cursor.advance(); + let finished = cursor.is_finished(); + if finished { + // Take the current cursor, leaving `None` in its place + let taken = self.cursors[stream_idx].take(); + if let Some(prev_cursors) = &mut self.prev_cursors { + prev_cursors[stream_idx] = taken; + } + } + return finished; + } + + // the entire stream is exhausted, so return true (poll won't help here anyway) + true + } + + /// Returns `true` if the cursor at index `a` is greater than at index `b`. + /// In an equality case, it compares the partition indices given. + #[inline] + fn is_gt(&self, a: usize, b: usize) -> bool { + match (&self.cursors[a], &self.cursors[b]) { + (None, _) => true, + (_, None) => false, + (Some(ac), Some(bc)) => ac.cmp(bc).then_with(|| a.cmp(&b)).is_gt(), + } + } + + #[inline] + fn is_poll_count_gt(&self, a: usize, b: usize) -> bool { + let poll_a = self.num_of_polled_with_same_value[a]; + let poll_b = self.num_of_polled_with_same_value[b]; + poll_a.cmp(&poll_b).then_with(|| a.cmp(&b)).is_gt() + } + + #[inline] + fn update_winner(&mut self, cmp_node: usize, winner: &mut usize, challenger: usize) { + self.loser_tree[cmp_node] = *winner; + *winner = challenger; + } + + /// Find the leaf node index in the loser tree for the given cursor index + /// + /// Note that this is not necessarily a leaf node in the tree, but it can + /// also be a half-node (a node with only one child). This happens when the + /// number of cursors/streams is not a power of two. Thus, the loser tree + /// will be unbalanced, but it will still work correctly. + /// + /// For example, with 5 streams, the loser tree will look like this: + /// + /// ```text + /// 0 (winner) + /// + /// 1 + /// / \ + /// 2 3 + /// / \ / \ + /// 4 | | | + /// / \ | | | + /// -+---+--+---+---+---- Below is not a part of loser tree + /// S3 S4 S0 S1 S2 + /// ``` + /// + /// S0, S1, ... S4 are the streams (read: stream at index 0, stream at + /// index 1, etc.) + /// + /// Zooming in at node 2 in the loser tree as an example, we can see that + /// it takes as input the next item at (S0) and the loser of (S3, S4). + #[inline] + fn lt_leaf_node_index(&self, cursor_index: usize) -> usize { + (self.cursors.len() + cursor_index) / 2 + } + + /// Find the parent node index for the given node index + #[inline] + fn lt_parent_node_index(&self, node_idx: usize) -> usize { + node_idx / 2 + } + + /// Attempts to initialize the loser tree with one value from each + /// non exhausted input, if possible + fn init_loser_tree(&mut self) { + // Init loser tree + self.loser_tree = vec![usize::MAX; self.cursors.len()]; + for i in 0..self.cursors.len() { + let mut winner = i; + let mut cmp_node = self.lt_leaf_node_index(i); + while cmp_node != 0 && self.loser_tree[cmp_node] != usize::MAX { + let challenger = self.loser_tree[cmp_node]; + if self.is_gt(winner, challenger) { + self.loser_tree[cmp_node] = winner; + winner = challenger; + } + + cmp_node = self.lt_parent_node_index(cmp_node); + } + self.loser_tree[cmp_node] = winner; + } + } + + /// Resets the poll count by incrementing the reset epoch. + fn reset_poll_counts(&mut self) { + self.current_reset_epoch += 1; + } + + /// Handles tie-breaking logic during the adjustment of the loser tree. + /// + /// When comparing elements from multiple partitions in the `update_loser_tree` process, a tie can occur + /// between the current winner and a challenger. This function is invoked when such a tie needs to be + /// resolved according to the round-robin tie-breaker mode. + /// + /// If round-robin tie-breaking is not active, it is enabled, and the poll counts for all elements are reset. + /// The function then compares the poll counts of the current winner and the challenger: + /// - If the winner remains at the top after the final comparison, it increments the winner's poll count. + /// - If the challenger has a lower poll count than the current winner, the challenger becomes the new winner. + /// - If the poll counts are equal but the challenger's index is smaller, the challenger is preferred. + /// + /// # Parameters + /// - `cmp_node`: The index of the comparison node in the loser tree where the tie-breaking is happening. + /// - `winner`: A mutable reference to the current winner, which may be updated based on the tie-breaking result. + /// - `challenger`: The index of the challenger being compared against the winner. + /// + /// This function ensures fair selection among elements with equal values when tie-breaking mode is enabled, + /// aiming to balance the polling across different partitions. + #[inline] + fn handle_tie(&mut self, cmp_node: usize, winner: &mut usize, challenger: usize) { + if !self.round_robin_tie_breaker_mode { + self.round_robin_tie_breaker_mode = true; + // Reset poll count for tie-breaker + self.reset_poll_counts(); + } + // Update poll count if the winner survives in the final match + if *winner == self.loser_tree[0] { + self.update_poll_count_on_the_same_value(*winner); + if self.is_poll_count_gt(*winner, challenger) { + self.update_winner(cmp_node, winner, challenger); + } + } else if challenger < *winner { + // If the winner doesn’t survive in the final match, it indicates that the original winner + // has moved up in value, so the challenger now becomes the new winner. + // This also means that we’re in a new round of the tie breaker, + // and the polls count is outdated (though not yet cleaned up). + // + // By the time we reach this code, both the new winner and the current challenger + // have the same value, and neither has an updated polls count. + // Therefore, we simply select the one with the smaller index. + self.update_winner(cmp_node, winner, challenger); + } + } + + /// Updates the loser tree to reflect the new winner after the previous winner is consumed. + /// This function adjusts the tree by comparing the current winner with challengers from + /// other partitions. + /// + /// If `enable_round_robin_tie_breaker` is true and a tie occurs at the final level, the + /// tie-breaker logic will be applied to ensure fair selection among equal elements. + fn update_loser_tree(&mut self) { + // Start with the current winner + let mut winner = self.loser_tree[0]; + + // Find the leaf node index of the winner in the loser tree. + let mut cmp_node = self.lt_leaf_node_index(winner); + + // Traverse up the tree to adjust comparisons until reaching the root. + while cmp_node > 1 { + let challenger = self.loser_tree[cmp_node]; + if self.is_gt(winner, challenger) { + self.update_winner(cmp_node, &mut winner, challenger); + } + cmp_node = self.lt_parent_node_index(cmp_node); + } + + if cmp_node == 1 { + let challenger = self.loser_tree[1]; + // If round-robin tie-breaker is enabled and we're at the final comparison (cmp_node == 1) + if self.round_robin_tie_breaker_enabled() { + match (&self.cursors[winner], &self.cursors[challenger]) { + (Some(ac), Some(bc)) => match ac.cmp(bc) { + std::cmp::Ordering::Equal => { + self.handle_tie(cmp_node, &mut winner, challenger); + } + std::cmp::Ordering::Greater => { + // Ends of tie breaker + self.round_robin_tie_breaker_mode = false; + self.update_winner(cmp_node, &mut winner, challenger); + } + std::cmp::Ordering::Less => { + // Ends of tie breaker + self.round_robin_tie_breaker_mode = false; + } + }, + (None, _) => { + // Challenger wins, update winner + // Ends of tie breaker + self.round_robin_tie_breaker_mode = false; + self.update_winner(cmp_node, &mut winner, challenger); + } + (_, None) => { + // Winner wins again + // Ends of tie breaker + self.round_robin_tie_breaker_mode = false; + } + } + } else if self.is_gt(winner, challenger) { + self.update_winner(cmp_node, &mut winner, challenger); + } + } + + self.loser_tree[0] = winner; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metrics::ExecutionPlanMetricsSet; + use crate::sorts::stream::PartitionedStream; + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryPool, UnboundedMemoryPool, + }; + use futures::TryStreamExt; + use std::cmp::Ordering; + + #[derive(Debug)] + struct EmptyPartitionedStream; + + impl PartitionedStream for EmptyPartitionedStream { + type Output = Result<(DummyValues, RecordBatch)>; + + fn partitions(&self) -> usize { + 1 + } + + fn poll_next( + &mut self, + _cx: &mut Context<'_>, + _stream_idx: usize, + ) -> Poll<Option<Self::Output>> { + Poll::Ready(None) + } + } + + #[derive(Debug)] + struct DummyValues; + + impl CursorValues for DummyValues { + fn len(&self) -> usize { + 0 + } + + fn eq(_l: &Self, _l_idx: usize, _r: &Self, _r_idx: usize) -> bool { + unreachable!("done-path test should not compare cursors") + } + + fn eq_to_previous(_cursor: &Self, _idx: usize) -> bool { + unreachable!("done-path test should not compare cursors") + } + + fn compare(_l: &Self, _l_idx: usize, _r: &Self, _r_idx: usize) -> Ordering { + unreachable!("done-path test should not compare cursors") + } + } + + #[tokio::test] + async fn test_done_drains_buffered_rows() { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("test").register(&pool); + let metrics = ExecutionPlanMetricsSet::new(); + + let mut stream = SortPreservingMergeStream::<DummyValues>::new( + Box::new(EmptyPartitionedStream), + Arc::clone(&schema), + BaselineMetrics::new(&metrics, 0), + 16, + Some(1), + reservation, + true, + ); + + // Simulate rows left buffered in `in_progress` (as happens when + // `build_record_batch` emits a partial batch on offset overflow). With + // an empty input stream the merge loop breaks immediately, so the only + // way these rows reach the consumer is the generator's final drain loop. + let batch = + RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]) + .unwrap(); + stream.in_progress.push_batch(0, batch).unwrap(); + stream.in_progress.push_row(0); + + // Drive the actual stream and confirm the buffered row is drained. + let batches: Vec<RecordBatch> = stream.into_stream().try_collect().await.unwrap(); + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/mod.rs b/native/vendor/datafusion-physical-plan/src/sorts/mod.rs new file mode 100644 index 0000000..ca8d4a4 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/mod.rs
@@ -0,0 +1,31 @@ +// 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 functionalities + +mod builder; +mod cursor; +mod merge; +mod multi_level_merge; +pub mod partial_sort; +pub mod partitioned_topk; +pub mod sort; +pub mod sort_preserving_merge; +mod stream; +pub mod streaming_merge; + +pub(crate) use stream::IncrementalSortIterator;
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/multi_level_merge.rs b/native/vendor/datafusion-physical-plan/src/sorts/multi_level_merge.rs new file mode 100644 index 0000000..3ec52cc --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/multi_level_merge.rs
@@ -0,0 +1,1100 @@ +// 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. + +//! Create a stream that do a multi level merge stream + +use crate::metrics::BaselineMetrics; +use crate::{EmptyRecordBatchStream, SpillManager}; +use arrow::array::RecordBatch; +use std::fmt::{Debug, Formatter}; +use std::mem; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, internal_err, resources_err}; +use datafusion_execution::memory_pool::MemoryReservation; + +use crate::sorts::builder::try_grow_reservation_to_at_least; +use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::TryStreamExt; +use futures::{Stream, StreamExt}; + +/// Merges a stream of sorted cursors and record batches into a single sorted stream +/// +/// This is a wrapper around [`SortPreservingMergeStream`](crate::sorts::merge::SortPreservingMergeStream) +/// that provide it the sorted streams/files to merge while making sure we can merge them in memory. +/// In case we can't merge all of them in a single pass we will spill the intermediate results to disk +/// and repeat the process. +/// +/// ## High level Algorithm +/// 1. Get the maximum amount of sorted in-memory streams and spill files we can merge with the available memory +/// 2. Sort them to a sorted stream +/// 3. Do we have more spill files to merge? +/// - Yes: write that sorted stream to a spill file, +/// add that spill file back to the spill files to merge and +/// repeat the process +/// +/// - No: return that sorted stream as the final output stream +/// +/// ```text +/// Initial State: Multiple sorted streams + spill files +/// ┌───────────┐ +/// │ Phase 1 │ +/// └───────────┘ +/// ┌──Can hold in memory─┐ +/// │ ┌──────────────┐ │ +/// │ │ In-memory │ +/// │ │sorted stream │──┼────────┐ +/// │ │ 1 │ │ │ +/// └──────────────┘ │ │ +/// │ ┌──────────────┐ │ │ +/// │ │ In-memory │ │ +/// │ │sorted stream │──┼────────┤ +/// │ │ 2 │ │ │ +/// └──────────────┘ │ │ +/// │ ┌──────────────┐ │ │ +/// │ │ In-memory │ │ +/// │ │sorted stream │──┼────────┤ +/// │ │ 3 │ │ │ +/// └──────────────┘ │ │ +/// │ ┌──────────────┐ │ │ ┌───────────┐ +/// │ │ Sorted Spill │ │ │ Phase 2 │ +/// │ │ file 1 │──┼────────┤ └───────────┘ +/// │ └──────────────┘ │ │ +/// ──── ──── ──── ──── ─┘ │ ┌──Can hold in memory─┐ +/// │ │ │ +/// ┌──────────────┐ │ │ ┌──────────────┐ +/// │ Sorted Spill │ │ │ │ Sorted Spill │ │ +/// │ file 2 │──────────────────────▶│ file 2 │──┼─────┐ +/// └──────────────┘ │ └──────────────┘ │ │ +/// ┌──────────────┐ │ │ ┌──────────────┐ │ │ +/// │ Sorted Spill │ │ │ │ Sorted Spill │ │ +/// │ file 3 │──────────────────────▶│ file 3 │──┼─────┤ +/// └──────────────┘ │ │ └──────────────┘ │ │ +/// ┌──────────────┐ │ ┌──────────────┐ │ │ +/// │ Sorted Spill │ │ │ │ Sorted Spill │ │ │ +/// │ file 4 │──────────────────────▶│ file 4 │────────┤ ┌───────────┐ +/// └──────────────┘ │ │ └──────────────┘ │ │ │ Phase 3 │ +/// │ │ │ │ └───────────┘ +/// │ ──── ──── ──── ──── ─┘ │ ┌──Can hold in memory─┐ +/// │ │ │ │ +/// ┌──────────────┐ │ ┌──────────────┐ │ │ ┌──────────────┐ +/// │ Sorted Spill │ │ │ Sorted Spill │ │ │ │ Sorted Spill │ │ +/// │ file 5 │──────────────────────▶│ file 5 │────────────────▶│ file 5 │───┼───┐ +/// └──────────────┘ │ └──────────────┘ │ │ └──────────────┘ │ │ +/// │ │ │ │ │ +/// │ ┌──────────────┐ │ │ ┌──────────────┐ │ +/// │ │ Sorted Spill │ │ │ │ Sorted Spill │ │ │ ┌── ─── ─── ─── ─── ─── ─── ──┐ +/// └──────────▶│ file 6 │────────────────▶│ file 6 │───┼───┼──────▶ Output Stream +/// └──────────────┘ │ │ └──────────────┘ │ │ └── ─── ─── ─── ─── ─── ─── ──┘ +/// │ │ │ │ +/// │ │ ┌──────────────┐ │ +/// │ │ │ Sorted Spill │ │ │ +/// └───────▶│ file 7 │───┼───┘ +/// │ └──────────────┘ │ +/// │ │ +/// └─ ──── ──── ──── ──── +/// ``` +/// +/// ## Memory Management Strategy +/// +/// This multi-level merge make sure that we can handle any amount of data to sort as long as +/// we have enough memory to merge at least 2 streams at a time, even when individual record +/// batches are skewed (very wide). +/// +/// 1. **Worst-Case Memory Reservation**: Reserves memory based on the largest +/// batch size encountered in each spill file to merge, ensuring sufficient memory is always +/// available during merge operations. +/// 2. **Adaptive Buffer Sizing**: Reduces buffer sizes when memory is constrained +/// 3. **Spill-to-Disk**: Spill to disk when we cannot merge all files in memory +/// 4. **Re-spilling Skewed Runs**: If even at the smallest read-buffer size we still cannot +/// reserve memory for the minimum of 2 streams - because a single run's largest batch is so +/// wide that two streams' worth of reservation exceeds the budget - the larger of the two +/// runs is re-spilled with each batch sliced in half. This shrinks its largest batch, +/// lowering the per-stream reservation, and the merge pass is retried. The re-spilled run +/// is tracked alongside a per-run batch-size limit equal to half the batch size it was +/// written with, so any later merge that includes it caps its output batch size to match - +/// otherwise the merged run could rebuild a full-size batch and reintroduce the skew. +/// Crucially the global merge batch size is *not* lowered, so re-spilling more than one run +/// does not compound the reduction. If a batch cannot be split any further (a single row +/// wider than the budget), the merge surfaces `ResourcesExhausted` instead of looping +/// forever. +pub(crate) struct MultiLevelMergeBuilder { + spill_manager: SpillManager, + schema: SchemaRef, + /// Sorted runs still to be merged. Each run is paired with the batch-size limit a + /// merge consuming it must cap its output at. Runs written at the full batch size + /// carry `batch_size`. A run re-spilled smaller to resolve skew carries its halved + /// limit (see [`Self::split_spill_file_in_half`]). Tracking it here keeps this limit + /// out of the public [`SortedSpillFile`], so no external caller has to set it. + sorted_spill_files: Vec<(SortedSpillFile, usize)>, + sorted_streams: Vec<SendableRecordBatchStream>, + expr: LexOrdering, + metrics: BaselineMetrics, + batch_size: usize, + reservation: MemoryReservation, + fetch: Option<usize>, + enable_round_robin_tie_breaker: bool, +} + +impl Debug for MultiLevelMergeBuilder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "MultiLevelMergeBuilder") + } +} + +impl MultiLevelMergeBuilder { + #[expect(clippy::too_many_arguments)] + pub(crate) fn new( + spill_manager: SpillManager, + schema: SchemaRef, + sorted_spill_files: Vec<SortedSpillFile>, + sorted_streams: Vec<SendableRecordBatchStream>, + expr: LexOrdering, + metrics: BaselineMetrics, + batch_size: usize, + reservation: MemoryReservation, + fetch: Option<usize>, + enable_round_robin_tie_breaker: bool, + ) -> Self { + Self { + spill_manager, + schema, + // Initial runs are written at the full batch size, so they impose no cap + // on later merges - record `batch_size` as their (unconstrained) limit. + sorted_spill_files: sorted_spill_files + .into_iter() + .map(|file| (file, batch_size)) + .collect(), + sorted_streams, + expr, + metrics, + batch_size, + reservation, + enable_round_robin_tie_breaker, + fetch, + } + } + + pub(crate) fn create_spillable_merge_stream(self) -> SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + futures::stream::once(self.create_stream()).try_flatten(), + )) + } + + async fn create_stream(mut self) -> Result<SendableRecordBatchStream> { + loop { + let (mut stream, batch_size_limit) = + match self.merge_sorted_runs_within_mem_limit()? { + MergeStep::Stream { + stream, + batch_size_limit, + } => (stream, batch_size_limit), + MergeStep::SplitThenRetry(index) => { + // Couldn't reserve memory for the minimum of 2 streams. Re-spill + // the larger of the two we're trying to merge with half its batch + // size so its largest batch shrinks, lowering the per-stream + // reservation, then retry. Makes the merge resilient to skewed + // (very wide) rows. + self.split_spill_file_in_half(index).await?; + continue; + } + }; + + // TODO - add a threshold for number of files to disk even if empty and reading from disk so + // we can avoid the memory reservation + + // If no spill files are left, we can return the stream as this is the last sorted run + // TODO - We can write to disk before reading it back to avoid having multiple streams in memory + if self.sorted_spill_files.is_empty() { + assert!( + self.sorted_streams.is_empty(), + "We should not have any sorted streams left" + ); + + return Ok(stream); + } + + // Need to sort to a spill file + let Some((spill_file, max_record_batch_memory)) = self + .spill_manager + .spill_record_batch_stream_and_return_max_batch_memory( + &mut stream, + "MultiLevelMergeBuilder intermediate spill", + ) + .await? + else { + continue; + }; + + // Add the spill file paired with the batch-size limit of the merge that + // produced it: if that merge consumed a shrunk (skew-resolved) run, its + // output was capped and this intermediate run is likewise capped, so a + // later pass that re-merges it won't rebuild an oversized batch. + self.sorted_spill_files.push(( + SortedSpillFile { + file: spill_file, + max_record_batch_memory, + }, + batch_size_limit, + )); + } + } + + /// This tries to create a stream that merges the most sorted streams and sorted spill files + /// as possible within the memory limit. + fn merge_sorted_runs_within_mem_limit(&mut self) -> Result<MergeStep> { + match (self.sorted_spill_files.len(), self.sorted_streams.len()) { + // No data so empty batch + (0, 0) => { + let empty_stream = + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); + Ok(MergeStep::Stream { + stream: self.observe_output(empty_stream), + batch_size_limit: self.batch_size, + }) + } + + // Only in-memory stream, return that + (0, 1) => { + let output_stream = self.sorted_streams.remove(0); + Ok(MergeStep::Stream { + stream: self.observe_output(output_stream), + batch_size_limit: self.batch_size, + }) + } + + // Only single sorted spill file so return it + (1, 0) => { + let (spill_file, batch_size) = self.sorted_spill_files.remove(0); + + // Not reserving any memory for this disk as we are not holding it in memory + let output_stream = self + .spill_manager + .read_spill_as_stream(spill_file.file, None)?; + + Ok(MergeStep::Stream { + stream: self.observe_output(output_stream), + batch_size_limit: batch_size, + }) + } + + // Only in memory streams, so merge them all in a single pass. In-memory + // runs are never shrunk for skew, so this merge runs at the full batch + // size and its output carries no limit. + (0, _) => { + let sorted_stream = mem::take(&mut self.sorted_streams); + // No need to wrap with observed stream since merge sort will update the observed metrics + Ok(MergeStep::Stream { + stream: self.create_new_merge_sort( + sorted_stream, + // If we have no sorted spill files left, this is the last run + true, + true, + self.batch_size, + )?, + batch_size_limit: self.batch_size, + }) + } + + // Need to merge multiple streams + (_, _) => { + // Transfer any pre-reserved bytes (from sort_spill_reservation_bytes) + // to the merge memory reservation. This prevents starvation when + // concurrent sort partitions compete for pool memory: the pre-reserved + // bytes cover spill file buffer reservations without additional pool + // allocation. + let mut memory_reservation = self.reservation.take(); + + // Compute the minimum before taking the in-memory streams so that, if we + // need to re-spill and retry, `self.sorted_streams` is left untouched. + let minimum_number_of_required_streams = + 2_usize.saturating_sub(self.sorted_streams.len()); + + let (sorted_spill_files, buffer_size) = match self + .get_sorted_spill_files_to_merge( + 2, + // we must have at least 2 streams to merge + minimum_number_of_required_streams, + &mut memory_reservation, + )? { + SpillFilesToMerge::Ready(sorted_spill_files, buffer_size) => { + (sorted_spill_files, buffer_size) + } + // Not enough memory to seat 2 streams. Re-spill the blocking file + // smaller and retry. `get_sorted_spill_files_to_merge` already freed + // the reservation and `self.sorted_streams` is untouched, so the + // retry starts clean. + SpillFilesToMerge::SplitThenRetry(index) => { + return Ok(MergeStep::SplitThenRetry(index)); + } + }; + + // Don't account for existing streams memory + // as we are not holding the memory for them + let mut sorted_streams = mem::take(&mut self.sorted_streams); + + let is_only_merging_memory_streams = sorted_spill_files.is_empty(); + + // If no spill files were selected (e.g. all too large for + // available memory but enough in-memory streams exist), + // return the pre-reserved bytes to self.reservation so + // create_new_merge_sort can transfer them to the merge + // stream's BatchBuilder. + if is_only_merging_memory_streams { + mem::swap(&mut self.reservation, &mut memory_reservation); + } + + // Cap the merge output at the smallest limit among the runs we're + // about to merge. Runs that were shrunk for skew carry a smaller limit, + // if none do, every run carries `self.batch_size` and the merge runs at + // the full batch size. The output stream is tagged with the same limit + // (see the `MergeStep::Stream` returns below) so a re-spilled + // intermediate run stays shrunk and won't rebuild an oversized batch on + // a later pass. + let mut output_batch_size = self.batch_size; + for (spill, batch_size_limit) in sorted_spill_files { + let stream = self + .spill_manager + .clone() + .with_batch_read_buffer_capacity(buffer_size) + .read_spill_as_stream( + spill.file, + Some(spill.max_record_batch_memory), + )?; + output_batch_size = output_batch_size.min(batch_size_limit); + sorted_streams.push(stream); + } + let merge_sort_stream = self.create_new_merge_sort( + sorted_streams, + // If we have no sorted spill files left, this is the last run + self.sorted_spill_files.is_empty(), + is_only_merging_memory_streams, + output_batch_size, + )?; + + // If we're only merging memory streams, we don't need to attach the memory reservation + // as it's empty + if is_only_merging_memory_streams { + assert_eq!( + memory_reservation.size(), + 0, + "when only merging memory streams, we should not have any memory reservation and let the merge sort handle the memory" + ); + + Ok(MergeStep::Stream { + stream: merge_sort_stream, + batch_size_limit: output_batch_size, + }) + } else { + // Attach the memory reservation to the stream to make sure we have enough memory + // throughout the merge process as we bypassed the memory pool for the merge sort stream + Ok(MergeStep::Stream { + stream: Box::pin(StreamAttachedReservation::new( + merge_sort_stream, + memory_reservation, + )), + batch_size_limit: output_batch_size, + }) + } + } + } + } + + fn create_new_merge_sort( + &mut self, + streams: Vec<SendableRecordBatchStream>, + is_output: bool, + all_in_memory: bool, + output_batch_size: usize, + ) -> Result<SendableRecordBatchStream> { + let mut builder = StreamingMergeBuilder::new() + .with_schema(Arc::clone(&self.schema)) + .with_expressions(&self.expr) + .with_batch_size(output_batch_size) + .with_fetch(self.fetch) + .with_metrics(if is_output { + // Only add the metrics to the last run + self.metrics.clone() + } else { + self.metrics.intermediate() + }) + .with_round_robin_tie_breaker(self.enable_round_robin_tie_breaker) + .with_streams(streams); + + if !all_in_memory { + // Don't track memory used by this stream as we reserve that memory by worst case sceneries + // (reserving memory for the biggest batch in each stream) + // TODO - avoid this hack as this can be broken easily when `SortPreservingMergeStream` + // changes the implementation to use more/less memory + builder = builder.with_bypass_mempool(); + } else { + // If we are only merging in-memory streams, we need to use the memory reservation + // because we don't know the maximum size of the batches in the streams. + // Use take() to transfer any pre-reserved bytes so the merge can use them + // as its initial budget without additional pool allocation. + builder = builder.with_reservation(self.reservation.take()); + } + + builder.build() + } + + /// Return the sorted spill files to use for the next phase, and the buffer size + /// This will try to get as many spill files as possible to merge, and if we don't have enough streams + /// it will try to reduce the buffer size until we have enough streams to merge + /// otherwise it will return an error + fn get_sorted_spill_files_to_merge( + &mut self, + buffer_len: usize, + minimum_number_of_required_streams: usize, + reservation: &mut MemoryReservation, + ) -> Result<SpillFilesToMerge> { + assert_ne!(buffer_len, 0, "Buffer length must be greater than 0"); + let mut number_of_spills_to_read_for_current_phase = 0; + let configured_fan_in = self + .spill_manager + .env() + .disk_manager + .max_spill_merge_fan_in(); + let max_spill_files = effective_spill_merge_fan_in(configured_fan_in); + // Track total memory needed for spill file buffers. When the + // reservation has pre-reserved bytes (from sort_spill_reservation_bytes), + // those bytes cover the first N spill files without additional pool + // allocation, preventing starvation under memory pressure. + let mut total_needed: usize = 0; + + for (spill, _) in &self.sorted_spill_files { + if number_of_spills_to_read_for_current_phase >= max_spill_files { + break; + } + + let per_spill = get_reserved_bytes_for_record_batch_size( + spill.max_record_batch_memory, + // Size will be the same as the sliced size, bc it is a spilled batch. + spill.max_record_batch_memory, + ) * buffer_len; + total_needed += per_spill; + + // For memory pools that are not shared this is good, for other + // this is not and there should be some upper limit to memory + // reservation so we won't starve the system. + match try_grow_reservation_to_at_least(reservation, total_needed) { + Ok(_) => { + number_of_spills_to_read_for_current_phase += 1; + } + // If we can't grow the reservation, we need to stop + Err(err) => { + // We must have at least 2 streams to merge, so if we don't have enough memory + // fail + if minimum_number_of_required_streams + > number_of_spills_to_read_for_current_phase + { + // Free the memory we reserved for this merge as we either try again or fail + reservation.free(); + if buffer_len > 1 { + // Try again with smaller buffer size, it will be slower but at least we can merge + return self.get_sorted_spill_files_to_merge( + buffer_len - 1, + minimum_number_of_required_streams, + reservation, + ); + } + + // buffer_len == 1 and we still can't seat the minimum of 2 streams. + if number_of_spills_to_read_for_current_phase == 0 { + // We couldn't even reserve a single stream - one record batch + // is larger than the whole merge budget. That's the lone-batch + // case, not the 2-stream merge skew we rescue here - surface it. + return Err(err); + } + + // We seated one stream (index 0) but not the second (index 1, the + // batch that just failed to reserve). Those are by definition the + // only two streams we are trying to merge, so re-spill the larger + // of them with a smaller batch size and retry, the smaller max + // batch lowers the per-stream reservation enough to seat both. + let split_index = usize::from( + self.sorted_spill_files[1].0.max_record_batch_memory + > self.sorted_spill_files[0].0.max_record_batch_memory, + ); + return Ok(SpillFilesToMerge::SplitThenRetry(split_index)); + } + + // We reached the maximum amount of memory we can use + // for this merge + break; + } + } + } + + let spills = self + .sorted_spill_files + .drain(..number_of_spills_to_read_for_current_phase) + .collect::<Vec<_>>(); + + Ok(SpillFilesToMerge::Ready(spills, buffer_len)) + } + + /// Re-spill the spill file at `index` with half its batch size, putting it back + /// at the same position. We read the file back and re-spill it through the normal + /// spill API (which owns batch layout), slicing every batch in two, which halves + /// the largest written batch and so lowers the per-stream merge reservation enough + /// for the next attempt to seat both streams. One stream's worth of memory is + /// reserved for the duration and freed afterwards. Makes the merge resilient to skew. + /// + /// Instead of halving the *global* merge batch size (which would compound when more + /// than one run is re-spilled), the shrunk run records its own smaller batch-size + /// limit (tracked alongside the run in `sorted_spill_files`), so only merges that + /// actually consume it pay the reduced batch size. + async fn split_spill_file_in_half(&mut self, index: usize) -> Result<()> { + log::debug!( + "2 spilled streams could not be loaded into memory for merge \ + (requires 2x of the largest batch from both), re-spilling the larger of the two with half \ + the batch size to reduce memory needs for the next merge attempt. the shrunk run carries \ + a halved batch-size limit so only merges consuming it use the smaller batch size" + ); + + // Extract the target in O(1) instead of `remove(index)`, which would shift + // every following spill file. Swap it to the back and pop it; the matching + // swap after re-spilling restores the original order, so the vec ends up + // exactly as it started, just with the target file shrunk. + // `old_batch_size` is the batch size this run was written with (the full merge + // batch size unless it was already shrunk once). Halving it caps the next merge + // that reads this run so the merged output can't rebuild a full-size batch. + let last = self.sorted_spill_files.len() - 1; + self.sorted_spill_files.swap(index, last); + let (target, old_batch_size) = self + .sorted_spill_files + .pop() + .expect("index is in bounds, so the vec is non-empty"); + let old_max = target.max_record_batch_memory; + + // Reserve enough to hold a single stream of this file while we re-spill it. + let reservation = self.reservation.new_empty(); + reservation + .try_grow(get_reserved_bytes_for_record_batch_size(old_max, old_max))?; + + let source = self + .spill_manager + .read_spill_as_stream(target.file, Some(old_max))?; + // Re-spill with half the batch size: slice every batch in two. The spill + // writer owns the batch layout, we only change how many rows per batch. + let mut halved: SendableRecordBatchStream = + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + source.flat_map(|batch| { + futures::stream::iter(match batch { + Ok(batch) => split_batch_in_half(batch) + .into_iter() + .map(Ok) + .collect::<Vec<_>>(), + Err(e) => vec![Err(e)], + }) + }), + )); + + let result = self + .spill_manager + .spill_record_batch_stream_and_return_max_batch_memory( + &mut halved, + "MultiLevelMergeBuilder split skewed spill", + ) + .await?; + + reservation.free(); + + let Some((file, new_max)) = result else { + return internal_err!("re-spilling a skewed spill file produced no data"); + }; + + // If halving could not reduce the largest batch (e.g. a single row that is + // itself wider than the budget), there is nothing more we can do - surface + // the out-of-memory condition instead of looping forever. + if new_max >= old_max { + return resources_err!( + "Cannot merge sorted runs: a single record batch of {old_max} bytes \ + exceeds the available merge memory and cannot be split further" + ); + } + + // Record the halved batch size as a *per-run* limit rather than lowering the + // global batch size. Merges that don't touch this run keep the full batch + // size. a merge that reads it caps its output at this limit so the merged run + // can't rebuild a full-size batch and reintroduce the skew. + let new_batch_size_limit = (old_batch_size / 2).max(1); + + // Push the re-spilled (smaller) file and swap it back into `index`, undoing + // the swap-to-back above so the order is preserved. + self.sorted_spill_files.push(( + SortedSpillFile { + file, + max_record_batch_memory: new_max, + }, + new_batch_size_limit, + )); + let last = self.sorted_spill_files.len() - 1; + self.sorted_spill_files.swap(index, last); + + Ok(()) + } + + fn observe_output( + &self, + stream: SendableRecordBatchStream, + ) -> SendableRecordBatchStream { + Box::pin(ObservedStream::new(stream, self.metrics.clone(), None)) + } +} + +/// Outcome of trying to reserve memory for one multi-level merge pass. +enum SpillFilesToMerge { + /// Enough memory: the spill files to read this pass (each paired with its + /// batch-size limit) and the read-ahead buffer size. + Ready(Vec<(SortedSpillFile, usize)>, usize), + /// Could not seat the minimum of 2 streams. Re-spill the spill file at this index + /// with a smaller (halved) batch size, then retry the pass. + SplitThenRetry(usize), +} + +/// What one iteration of the multi-level merge loop should do next. +enum MergeStep { + /// A merged stream is ready to be consumed (and possibly spilled back). + Stream { + stream: SendableRecordBatchStream, + /// The batch-size limit to stamp on the run if this stream is re-spilled as an + /// intermediate result: the batch size its merge ran at. It equals the full + /// merge batch size unless the merge consumed a skew-resolved run, in which + /// case it is that run's smaller limit so the re-spilled result stays capped + /// and can't rebuild an oversized batch. + batch_size_limit: usize, + }, + /// Re-spill the spill file at this index smaller, then retry the merge step. + SplitThenRetry(usize), +} + +/// Slice `batch` into two row-halves so a re-spill writes batches half the size. +fn split_batch_in_half(batch: RecordBatch) -> Vec<RecordBatch> { + let num_rows = batch.num_rows(); + if num_rows <= 1 { + return vec![batch]; + } + let mid = num_rows / 2; + vec![batch.slice(0, mid), batch.slice(mid, num_rows - mid)] +} + +fn effective_spill_merge_fan_in(configured_fan_in: usize) -> usize { + if configured_fan_in == 0 { + usize::MAX + } else { + configured_fan_in.max(2) + } +} + +struct StreamAttachedReservation { + stream: SendableRecordBatchStream, + reservation: MemoryReservation, +} + +impl StreamAttachedReservation { + fn new(stream: SendableRecordBatchStream, reservation: MemoryReservation) -> Self { + Self { + stream, + reservation, + } + } +} + +impl Stream for StreamAttachedReservation { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let res = self.stream.poll_next_unpin(cx); + + match res { + Poll::Ready(res) => { + match res { + Some(Ok(batch)) => Poll::Ready(Some(Ok(batch))), + Some(Err(err)) => { + // Had an error so drop the data + self.reservation.free(); + Poll::Ready(Some(Err(err))) + } + None => { + // Stream is done so free the memory + self.reservation.free(); + + Poll::Ready(None) + } + } + } + Poll::Pending => Poll::Pending, + } + } +} + +impl RecordBatchStream for StreamAttachedReservation { + fn schema(&self) -> SchemaRef { + self.stream.schema() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::expressions::PhysicalSortExpr; + use arrow::array::{AsArray, Int64Array}; + use arrow::compute::concat_batches; + use arrow::datatypes::{DataType, Field, Int64Type, Schema}; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; + use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + use datafusion_physical_expr::expressions::{Column, col}; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)])) + } + + fn build_spill_manager(env: &Arc<RuntimeEnv>, schema: &SchemaRef) -> SpillManager { + SpillManager::new( + Arc::clone(env), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(schema), + ) + } + + /// Spill `values` (which must already be sorted) as a single sorted run and + /// return it as a `SortedSpillFile` carrying its recorded largest-batch memory. + fn make_sorted_spill_file( + spill_manager: &SpillManager, + schema: &SchemaRef, + values: Vec<i64>, + ) -> SortedSpillFile { + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(Int64Array::from(values))], + ) + .unwrap(); + let batches: Vec<Result<RecordBatch>> = vec![Ok(batch)]; + let (file, max_record_batch_memory) = spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + batches.into_iter(), + "test input run", + ) + .unwrap() + .expect("spill should produce a file"); + SortedSpillFile { + file, + max_record_batch_memory, + } + } + + fn build_merge_builder( + spill_manager: SpillManager, + schema: SchemaRef, + sorted_spill_files: Vec<SortedSpillFile>, + pool: &Arc<dyn MemoryPool>, + batch_size: usize, + ) -> MultiLevelMergeBuilder { + let reservation = MemoryConsumer::new("test merge").register(pool); + let expr: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(); + MultiLevelMergeBuilder::new( + spill_manager, + schema, + sorted_spill_files, + vec![], + expr, + BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + batch_size, + reservation, + None, + false, + ) + } + + /// Two sorted runs whose largest batches are too big to both + /// be seated in the merge budget at once are re-spilled (halved) until they + /// fit, and the merge then completes with fully sorted, complete output. + #[tokio::test] + async fn skewed_runs_are_respilled_so_the_merge_fits() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + + // Seating two streams needs ~4*m (2*m each), which does NOT fit, but the + // budget is large enough once a run is halved. The rescue keeps halving + // the blocking run until two streams fit (here, after one halving). + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); + + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + 8192, + ); + let stream = builder.create_spillable_merge_stream(); + let batches: Vec<RecordBatch> = stream.try_collect().await?; + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, + (2 * n) as usize, + "the merge must emit every input row" + ); + + let merged = concat_batches(&schema, &batches)?; + let col = merged.column(0).as_primitive::<Int64Type>(); + for i in 1..col.len() { + assert!( + col.value(i - 1) <= col.value(i), + "merge output must be sorted: {} > {} at {i}", + col.value(i - 1), + col.value(i), + ); + } + + Ok(()) + } + + /// Tests the `new_max >= old_max` guard: a single-row run cannot be split + /// any smaller, so re-spilling it does not shrink the largest batch and the + /// rescue surfaces `ResourcesExhausted` rather than looping forever. + #[tokio::test] + async fn respilling_an_unsplittable_run_surfaces_resources_exhausted() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + // A one-row run: `split_batch_in_half` returns it unchanged, so the + // re-spilled file's largest batch cannot drop below the original. + let f0 = make_sorted_spill_file(&spill_manager, &schema, vec![42]); + + // Ample budget so the only possible failure is the un-splittable guard, + // not the single-stream reservation itself. + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let mut builder = + build_merge_builder(spill_manager, schema, vec![f0], &pool, 1024); + + let err = builder + .split_spill_file_in_half(0) + .await + .expect_err("re-spilling a one-row run cannot shrink it"); + assert!( + err.to_string().contains("cannot be split further"), + "expected the un-splittable guard error, got: {err}" + ); + + Ok(()) + } + + /// Proves the re-spill also halves the merge output batch size: after one + /// re-spill the merged run is emitted in 4096-row batches (not the original + /// 8192), so it cannot rebuild a full-size batch and reintroduce the skew. + #[tokio::test] + async fn respill_halves_the_merge_output_batch_size() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + + // 3.5*m forces exactly one re-spill (split one run, then both fit), which + // halves the merge output batch size. + let initial_batch_size = 8192; + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); + + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + initial_batch_size, + ); + let stream = builder.create_spillable_merge_stream(); + let batches: Vec<RecordBatch> = stream.try_collect().await?; + + // All rows are still present. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, (2 * n) as usize); + + // The largest emitted batch is the halved size, not the original 8192: the + // shrunk run carries a halved batch-size limit, and the final pass consumes + // it, so the merge output is capped there. Without the per-run limit the merge + // would rebuild 8192-row batches. + let expected_batch_size = initial_batch_size / 2; + let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); + assert_eq!( + max_batch_rows, expected_batch_size, + "after one re-spill the merge must emit {expected_batch_size}-row \ + batches, got a largest batch of {max_batch_rows} rows" + ); + + Ok(()) + } + + /// Same as [`respill_halves_the_merge_output_batch_size`], but under a budget tight + /// enough that *both* runs must be re-spilled before the merge fits - the scenario + /// where the batch-size reduction could compound. Because the reduction is tracked + /// per-run (each run capped at half) rather than by halving the global batch size on + /// every split, the merged output is emitted in 4096-row batches - half, not a + /// quarter. A global-halving implementation would have halved once per re-spill and + /// emitted 2048-row batches. + #[tokio::test] + async fn respilling_two_skewed_runs_halves_the_output_without_compounding() + -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + + // 2.5*m is tight enough that even after halving one run the two still don't + // fit, so *both* runs are re-spilled once before the merge succeeds. (3.5*m, + // as in the single-split test, would let the pair fit after one split.) This + // is exactly the scenario where a compounding, global-halving implementation + // would drive the output batch size down to a quarter. + let initial_batch_size = 8192; + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 5 / 2)); + + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + initial_batch_size, + ); + let stream = builder.create_spillable_merge_stream(); + let batches: Vec<RecordBatch> = stream.try_collect().await?; + + // All rows are still present. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, (2 * n) as usize); + + // Each run was re-spilled once, so each is capped at half the original batch + // size and the merge caps its output at that half - NOT a quarter. A global + // halving-per-split implementation would have emitted 2048-row batches here. + let expected_batch_size = initial_batch_size / 2; + let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); + assert_eq!( + max_batch_rows, expected_batch_size, + "two re-spills must halve (not quarter) the output: expected \ + {expected_batch_size}-row batches, got a largest batch of \ + {max_batch_rows} rows" + ); + + Ok(()) + } + + #[test] + fn spill_merge_fan_in_is_unlimited_by_default() { + assert_eq!(effective_spill_merge_fan_in(0), usize::MAX); + } + + #[test] + fn spill_merge_fan_in_preserves_merge_progress() { + assert_eq!(effective_spill_merge_fan_in(1), 2); + assert_eq!(effective_spill_merge_fan_in(2), 2); + assert_eq!(effective_spill_merge_fan_in(8), 8); + } + + #[test] + fn spill_merge_phase_respects_configured_fan_in() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let runtime = RuntimeEnvBuilder::new() + .with_max_spill_merge_fan_in(2) + .build_arc()?; + let spill_manager = SpillManager::new( + Arc::clone(&runtime), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(&schema), + ); + let sorted_spill_files = (0..4) + .map(|idx| { + Ok(SortedSpillFile { + file: runtime + .disk_manager + .create_tmp_file(&format!("spill fan-in test {idx}"))?, + max_record_batch_memory: 1, + }) + }) + .collect::<Result<Vec<_>>>()?; + let expr = LexOrdering::new([PhysicalSortExpr::new_default(col("a", &schema)?)]) + .unwrap(); + let reservation = + MemoryConsumer::new("spill_merge_phase_respects_configured_fan_in") + .register(&runtime.memory_pool); + let metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut builder = MultiLevelMergeBuilder::new( + spill_manager, + schema, + sorted_spill_files, + vec![], + expr, + metrics, + 1024, + reservation, + None, + false, + ); + let mut merge_reservation = MemoryConsumer::new("spill_merge_fan_in_phase") + .register(&runtime.memory_pool); + + let (spills, buffer_len) = match builder.get_sorted_spill_files_to_merge( + 1, + 2, + &mut merge_reservation, + )? { + SpillFilesToMerge::Ready(spills, buffer_len) => (spills, buffer_len), + SpillFilesToMerge::SplitThenRetry(index) => { + panic!("expected ready spill files, got retry for index {index}") + } + }; + + assert_eq!(spills.len(), 2); + assert_eq!(buffer_len, 1); + assert_eq!(builder.sorted_spill_files.len(), 2); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/partial_sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/partial_sort.rs new file mode 100644 index 0000000..478ac14 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/partial_sort.rs
@@ -0,0 +1,1347 @@ +// 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. + +//! Partial Sort deals with input data that partially +//! satisfies the required sort order. Such an input data can be +//! partitioned into segments where each segment already has the +//! required information for lexicographic sorting so sorting +//! can be done without loading the entire dataset. +//! +//! Consider a sort plan having an input with ordering `a ASC, b ASC` +//! +//! ```text +//! +---+---+---+ +//! | a | b | d | +//! +---+---+---+ +//! | 0 | 0 | 3 | +//! | 0 | 0 | 2 | +//! | 0 | 1 | 1 | +//! | 0 | 2 | 0 | +//! +---+---+---+ +//! ``` +//! +//! and required ordering for the plan is `a ASC, b ASC, d ASC`. +//! The first 3 rows(segment) can be sorted as the segment already +//! has the required information for the sort, but the last row +//! requires further information as the input can continue with a +//! batch with a starting row where a and b does not change as below +//! +//! ```text +//! +---+---+---+ +//! | a | b | d | +//! +---+---+---+ +//! | 0 | 2 | 4 | +//! +---+---+---+ +//! ``` +//! +//! The plan concats incoming data with such last rows of previous input +//! and continues partial sorting of the segments. + +use std::fmt::Debug; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::sorts::sort::sort_batch; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::EmptyRecordBatchStream; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, +}; + +use arrow::compute::concat_batches; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::evaluate_partition_ranges; +use datafusion_execution::{RecordBatchStream, TaskContext}; +use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; + +use futures::{Stream, StreamExt, ready}; +use log::trace; + +/// Sort execution plan for inputs that are already partially sorted. +/// +/// This operator takes input ordered by a prefix of the required ordering, and +/// produces output ordered by the required ordering, emitting rows sooner +/// (streaming) and using less peak memory than [`SortExec`] which must buffer +/// all rows before producing any output. +/// +/// [`PartialSortExec`] relies on the property that rows with the same sort +/// prefix are contiguous, so it can sort one prefix group at a time, emitting +/// completed groups without reading (and buffering) the entire input. +/// +/// For example, if the required output is `(a, b, c)`, but the input is only +/// ordered by `(a, b)`, `PartialSortExec` sorts only within each `(a, b)` +/// group to produce output ordered by `(a, b, c)`. +/// +/// ```text +/// input ordered by a, b output ordered by a, b, c +/// +/// +---+---+---+ +---+---+---+ +/// | a | b | c | | a | b | c | +/// +---+---+---+ +---+---+---+ +/// | 0 | 0 | 3 | -- new group --> | 0 | 0 | 1 | +/// | 0 | 0 | 2 | | 0 | 0 | 2 | +/// | 0 | 0 | 1 | | 0 | 0 | 3 | +/// | 0 | 1 | 1 | -- new group --> | 0 | 1 | 1 | +/// | 0 | 2 | 4 | -- new group --> | 0 | 2 | 0 | +/// | 0 | 2 | 0 | | 0 | 2 | 4 | +/// | 1 | 0 | 5 | -- new group --> | 1 | 0 | 5 | +/// +---+---+---+ +---+---+---+ +/// ``` +/// +/// # Buffering and Emitting Rows +/// +/// [`PartialSortExec`] buffers rows only until it can *prove* a prefix group +/// will never be seen again, then sorts and emits buffered rows. A group is +/// guaranteed to never be seen again once a row with a *different* prefix +/// value arrives. This relies on the input's existing ordering guarantees. +/// +/// Using the example from above, rows accumulate in the in-memory buffer in +/// batches. As long as the `(a, b)` prefix keeps repeating, more rows are +/// buffered. +/// +/// ```text +/// Buffer +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 0 | 0 | 3 | +/// | 0 | 0 | 2 | +/// | 0 | 0 | 1 | +/// +---+---+---+ +/// ``` +/// +/// Once a batch arrives that contains a new `(a, b)` prefix, e.g. `(0, 2)`: +/// every buffered row for previous prefixes may be emitted: +/// +/// ```text +/// Buffer +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 0 | 0 | 3 | +/// | 0 | 0 | 2 | +/// | 0 | 0 | 1 | +/// | 0 | 1 | 1 | <-- first row of new batch, new prefix +/// | 0 | 2 | 4 | <-- new prefix +/// | 0 | 2 | 0 | +/// | 1 | 0 | 5 | <-- last row of new batch, new prefix +/// +---+---+---+ +/// ``` +/// +/// Once known complete, the buffered rows are sorted by the full `(a, b, c)` +/// ordering and emitted as a [`RecordBatch`]; Any rows from the most recently +/// seen prefix remain buffered (as more rows with the same prefix may arrive in +/// future batches. +/// +/// ```text +/// Emitted <-- fully sorted on (a, b, c) +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 0 | 0 | 1 | <-- completed group +/// | 0 | 0 | 2 | +/// | 0 | 0 | 3 | +/// | 0 | 2 | 0 | <-- completed group +/// | 0 | 2 | 4 | +/// | 0 | 1 | 1 | <-- completed group +/// +---+---+---+ +/// +/// Buffer +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 1 | 0 | 5 | <-- (possibly) in progress group +/// +---+---+---+ +/// ``` +/// +/// [`SortExec`]: crate::sorts::sort::SortExec +#[derive(Debug, Clone)] +pub struct PartialSortExec { + /// Input schema + pub(crate) input: Arc<dyn ExecutionPlan>, + /// Sort expressions + expr: LexOrdering, + /// Length of continuous matching columns of input that satisfy + /// the required ordering for the sort + common_prefix_length: usize, + /// Containing all metrics set created during sort + metrics_set: ExecutionPlanMetricsSet, + /// Preserve partitions of input plan. If false, the input partitions + /// will be sorted and merged into a single output partition. + preserve_partitioning: bool, + /// Fetch highest/lowest n results + fetch: Option<usize>, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl PartialSortExec { + /// Create a new partial sort execution plan + pub fn new( + expr: LexOrdering, + input: Arc<dyn ExecutionPlan>, + common_prefix_length: usize, + ) -> Self { + debug_assert!(common_prefix_length > 0); + let preserve_partitioning = false; + let cache = Self::compute_properties(&input, expr.clone(), preserve_partitioning) + .unwrap(); + Self { + input, + expr, + common_prefix_length, + metrics_set: ExecutionPlanMetricsSet::new(), + preserve_partitioning, + fetch: None, + cache: Arc::new(cache), + } + } + + /// Whether this `PartialSortExec` preserves partitioning of the children + pub fn preserve_partitioning(&self) -> bool { + self.preserve_partitioning + } + + /// Specify the partitioning behavior of this partial sort exec + /// + /// If `preserve_partitioning` is true, sorts each partition + /// individually, producing one sorted stream for each input partition. + /// + /// If `preserve_partitioning` is false, sorts and merges all + /// input partitions producing a single, sorted partition. + pub fn with_preserve_partitioning(mut self, preserve_partitioning: bool) -> Self { + self.preserve_partitioning = preserve_partitioning; + Arc::make_mut(&mut self.cache).partitioning = + Self::output_partitioning_helper(&self.input, self.preserve_partitioning); + self + } + + /// Modify how many rows to include in the result + /// + /// If None, then all rows will be returned, in sorted order. + /// If Some, then only the top `fetch` rows will be returned. + /// This can reduce the memory pressure required by the sort + /// operation since rows that are not going to be included + /// can be dropped. + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + /// Input schema + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Sort expressions + pub fn expr(&self) -> &LexOrdering { + &self.expr + } + + /// If `Some(fetch)`, limits output to only the first "fetch" items + pub fn fetch(&self) -> Option<usize> { + self.fetch + } + + /// Common prefix length + pub fn common_prefix_length(&self) -> usize { + self.common_prefix_length + } + + fn output_partitioning_helper( + input: &Arc<dyn ExecutionPlan>, + preserve_partitioning: bool, + ) -> Partitioning { + // Get output partitioning: + if preserve_partitioning { + input.output_partitioning().clone() + } else { + Partitioning::UnknownPartitioning(1) + } + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + sort_exprs: LexOrdering, + preserve_partitioning: bool, + ) -> Result<PlanProperties> { + // Calculate equivalence properties; i.e. reset the ordering equivalence + // class with the new ordering: + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.reorder(sort_exprs)?; + + // Get output partitioning: + let output_partitioning = + Self::output_partitioning_helper(input, preserve_partitioning); + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + input.pipeline_behavior(), + input.boundedness(), + )) + } +} + +impl DisplayAs for PartialSortExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let common_prefix_length = self.common_prefix_length; + match self.fetch { + Some(fetch) => { + write!( + f, + "PartialSortExec: TopK(fetch={fetch}), expr=[{}], common_prefix_length=[{common_prefix_length}]", + self.expr + ) + } + None => write!( + f, + "PartialSortExec: expr=[{}], common_prefix_length=[{common_prefix_length}]", + self.expr + ), + } + } + DisplayFormatType::TreeRender => match self.fetch { + Some(fetch) => { + writeln!(f, "{}", self.expr)?; + writeln!(f, "limit={fetch}") + } + None => { + writeln!(f, "{}", self.expr) + } + }, + } + } +} + +impl ExecutionPlan for PartialSortExec { + fn name(&self) -> &'static str { + "PartialSortExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { + vec![Distribution::UnspecifiedDistribution] + } else { + vec![Distribution::SinglePartition] + }) + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics_set: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_partial_sort = PartialSortExec::new( + self.expr.clone(), + Arc::clone(&children[0]), + self.common_prefix_length, + ) + .with_fetch(self.fetch) + .with_preserve_partitioning(self.preserve_partitioning); + + Ok(Arc::new(new_partial_sort)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start PartialSortExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + let input = self.input.execute(partition, Arc::clone(&context))?; + + trace!("End PartialSortExec's input.execute for partition: {partition}"); + + // Make sure common prefix length is larger than 0 + // Otherwise, we should use SortExec. + debug_assert!(self.common_prefix_length > 0); + + Ok(Box::pin(PartialSortStream { + input, + expr: self.expr.clone(), + common_prefix_length: self.common_prefix_length, + in_mem_batch: RecordBatch::new_empty(Arc::clone(&self.schema())), + fetch: self.fetch, + is_closed: false, + baseline_metrics: BaselineMetrics::new(&self.metrics_set, partition), + })) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics_set.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::clone(&input_stats[0])) + } +} + +struct PartialSortStream { + /// The input plan + input: SendableRecordBatchStream, + /// Sort expressions + expr: LexOrdering, + /// Length of prefix common to input ordering and required ordering of plan + /// should be more than 0 otherwise PartialSort is not applicable + common_prefix_length: usize, + /// Used as a buffer for part of the input not ready for sort + in_mem_batch: RecordBatch, + /// Fetch top N results + fetch: Option<usize>, + /// Whether the stream has finished returning all of its data or not + is_closed: bool, + /// Execution metrics + baseline_metrics: BaselineMetrics, +} + +impl Stream for PartialSortStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let poll = self.poll_next_inner(cx); + self.baseline_metrics.record_poll(poll) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + // we can't predict the size of incoming batches so re-use the size hint from the input + self.input.size_hint() + } +} + +impl RecordBatchStream for PartialSortStream { + fn schema(&self) -> SchemaRef { + self.input.schema() + } +} + +impl PartialSortStream { + fn poll_next_inner( + self: &mut Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + if self.is_closed { + return Poll::Ready(None); + } + loop { + // Check if we've already reached the fetch limit + if self.fetch == Some(0) { + self.is_closed = true; + // Release the input pipeline's resources. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + return Poll::Ready(None); + } + + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + // Merge new batch into in_mem_batch + self.in_mem_batch = concat_batches( + &self.schema(), + &[self.in_mem_batch.clone(), batch], + )?; + + // Check if we have a slice point, otherwise keep accumulating in `self.in_mem_batch`. + if let Some(slice_point) = self + .get_slice_point(self.common_prefix_length, &self.in_mem_batch)? + { + let sorted = self.in_mem_batch.slice(0, slice_point); + self.in_mem_batch = self.in_mem_batch.slice( + slice_point, + self.in_mem_batch.num_rows() - slice_point, + ); + let sorted_batch = sort_batch(&sorted, &self.expr, self.fetch)?; + if let Some(fetch) = self.fetch.as_mut() { + *fetch -= sorted_batch.num_rows(); + } + + if sorted_batch.num_rows() > 0 { + return Poll::Ready(Some(Ok(sorted_batch))); + } + } + } + Some(Err(e)) => return Poll::Ready(Some(Err(e))), + None => { + self.is_closed = true; + // Release the input pipeline's resources before sorting. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + // Once input is consumed, sort the rest of the inserted batches + let remaining_batch = self.sort_in_mem_batch()?; + return if remaining_batch.num_rows() > 0 { + Poll::Ready(Some(Ok(remaining_batch))) + } else { + Poll::Ready(None) + }; + } + }; + } + } + + /// Returns a sorted RecordBatch from in_mem_batches and clears in_mem_batches + /// + /// If fetch is specified for PartialSortStream `sort_in_mem_batch` will limit + /// the last RecordBatch returned and will mark the stream as closed + fn sort_in_mem_batch(self: &mut Pin<&mut Self>) -> Result<RecordBatch> { + let input_batch = self.in_mem_batch.clone(); + self.in_mem_batch = RecordBatch::new_empty(self.schema()); + let result = sort_batch(&input_batch, &self.expr, self.fetch)?; + if let Some(remaining_fetch) = self.fetch { + // remaining_fetch - result.num_rows() is always be >= 0 + // because result length of sort_batch with limit cannot be + // more than the requested limit + self.fetch = Some(remaining_fetch - result.num_rows()); + if remaining_fetch == result.num_rows() { + self.is_closed = true; + } + } + Ok(result) + } + + /// Return the end index of the second last partition if the batch + /// can be partitioned based on its already sorted columns + /// + /// Return None if the batch cannot be partitioned, which means the + /// batch does not have the information for a safe sort + fn get_slice_point( + &self, + common_prefix_len: usize, + batch: &RecordBatch, + ) -> Result<Option<usize>> { + let common_prefix_sort_keys = (0..common_prefix_len) + .map(|idx| self.expr[idx].evaluate_to_sort_column(batch)) + .collect::<Result<Vec<_>>>()?; + let partition_points = + evaluate_partition_ranges(batch.num_rows(), &common_prefix_sort_keys)?; + // If partition points are [0..100], [100..200], [200..300] + // we should return 200, which is the safest and furthest partition boundary + // Please note that we shouldn't return 300 (which is number of rows in the batch), + // because this boundary may change with new data. + if partition_points.len() >= 2 { + Ok(Some(partition_points[partition_points.len() - 2].end)) + } else { + Ok(None) + } + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use arrow::array::*; + use arrow::compute::SortOptions; + use arrow::datatypes::*; + use datafusion_common::test_util::batches_to_string; + use futures::FutureExt; + use insta::allow_duplicates; + use insta::assert_snapshot; + use itertools::Itertools; + + use crate::collect; + use crate::expressions::PhysicalSortExpr; + use crate::expressions::col; + use crate::sorts::sort::SortExec; + use crate::test; + use crate::test::TestMemoryExec; + use crate::test::assert_is_pending; + use crate::test::exec::{BlockingExec, assert_strong_count_converges_to_zero}; + + use super::*; + + #[tokio::test] + async fn test_partial_sort() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let source = test::build_table_scan_i32( + ("a", &vec![0, 0, 0, 1, 1, 1]), + ("b", &vec![1, 1, 2, 2, 3, 3]), + ("c", &vec![1, 0, 5, 4, 3, 2]), + ); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + + let partial_sort_exec = Arc::new(PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + Arc::clone(&source), + 2, + )); + + let result = collect(partial_sort_exec, Arc::clone(&task_ctx)).await?; + + assert_eq!(2, result.len()); + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 0 | 1 | 0 | + | 0 | 1 | 1 | + | 0 | 2 | 5 | + | 1 | 2 | 4 | + | 1 | 3 | 2 | + | 1 | 3 | 3 | + +---+---+---+ + "); + } + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partial_sort_with_fetch() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let source = test::build_table_scan_i32( + ("a", &vec![0, 0, 1, 1, 1]), + ("b", &vec![1, 2, 2, 3, 3]), + ("c", &vec![4, 3, 2, 1, 0]), + ); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + + for common_prefix_length in [1, 2] { + let partial_sort_exec = Arc::new( + PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + Arc::clone(&source), + common_prefix_length, + ) + .with_fetch(Some(4)), + ); + + let result = collect(partial_sort_exec, Arc::clone(&task_ctx)).await?; + + assert_eq!(2, result.len()); + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 0 | 1 | 4 | + | 0 | 2 | 3 | + | 1 | 2 | 2 | + | 1 | 3 | 0 | + +---+---+---+ + "); + } + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_partial_sort2() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let source_tables = [ + test::build_table_scan_i32( + ("a", &vec![0, 0, 0, 0, 1, 1, 1, 1]), + ("b", &vec![1, 1, 3, 3, 4, 4, 2, 2]), + ("c", &vec![7, 6, 5, 4, 3, 2, 1, 0]), + ), + test::build_table_scan_i32( + ("a", &vec![0, 0, 0, 0, 1, 1, 1, 1]), + ("b", &vec![1, 1, 3, 3, 2, 2, 4, 4]), + ("c", &vec![7, 6, 5, 4, 1, 0, 3, 2]), + ), + ]; + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + for (common_prefix_length, source) in + [(1, &source_tables[0]), (2, &source_tables[1])] + { + let partial_sort_exec = Arc::new(PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + Arc::clone(source), + common_prefix_length, + )); + + let result = collect(partial_sort_exec, Arc::clone(&task_ctx)).await?; + assert_eq!(2, result.len()); + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 0 | 1 | 6 | + | 0 | 1 | 7 | + | 0 | 3 | 4 | + | 0 | 3 | 5 | + | 1 | 2 | 0 | + | 1 | 2 | 1 | + | 1 | 4 | 2 | + | 1 | 4 | 3 | + +---+---+---+ + "); + } + } + Ok(()) + } + + fn prepare_partitioned_input() -> Arc<dyn ExecutionPlan> { + let batch1 = test::build_table_i32( + ("a", &vec![1; 100]), + ("b", &(0..100).rev().collect()), + ("c", &(0..100).rev().collect()), + ); + let batch2 = test::build_table_i32( + ("a", &[&vec![1; 25][..], &vec![2; 75][..]].concat()), + ("b", &(100..200).rev().collect()), + ("c", &(0..100).collect()), + ); + let batch3 = test::build_table_i32( + ("a", &[&vec![3; 50][..], &vec![4; 50][..]].concat()), + ("b", &(150..250).rev().collect()), + ("c", &(0..100).rev().collect()), + ); + let batch4 = test::build_table_i32( + ("a", &vec![4; 100]), + ("b", &(50..150).rev().collect()), + ("c", &(0..100).rev().collect()), + ); + let schema = batch1.schema(); + + TestMemoryExec::try_new_exec( + &[vec![batch1, batch2, batch3, batch4]], + Arc::clone(&schema), + None, + ) + .unwrap() as Arc<dyn ExecutionPlan> + } + + #[tokio::test] + async fn test_partitioned_input_partial_sort() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let mem_exec = prepare_partitioned_input(); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let option_desc = SortOptions { + descending: false, + nulls_first: false, + }; + let schema = mem_exec.schema(); + let partial_sort_exec = PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_desc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + Arc::clone(&mem_exec), + 1, + ); + let sort_exec = Arc::new(SortExec::new( + partial_sort_exec.expr.clone(), + Arc::clone(&partial_sort_exec.input), + )); + let result = collect(Arc::new(partial_sort_exec), Arc::clone(&task_ctx)).await?; + assert_eq!( + result.iter().map(|r| r.num_rows()).collect_vec(), + [125, 125, 150] + ); + + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + let partial_sort_result = concat_batches(&schema, &result).unwrap(); + let sort_result = collect(sort_exec, Arc::clone(&task_ctx)).await?; + assert_eq!(sort_result[0], partial_sort_result); + + Ok(()) + } + + #[tokio::test] + async fn test_partitioned_input_partial_sort_with_fetch() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let mem_exec = prepare_partitioned_input(); + let schema = mem_exec.schema(); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let option_desc = SortOptions { + descending: false, + nulls_first: false, + }; + for (fetch_size, expected_batch_num_rows) in [ + (Some(50), vec![50]), + (Some(120), vec![120]), + (Some(150), vec![125, 25]), + (Some(250), vec![125, 125]), + ] { + let partial_sort_exec = PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_desc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + Arc::clone(&mem_exec), + 1, + ) + .with_fetch(fetch_size); + + let sort_exec = Arc::new( + SortExec::new( + partial_sort_exec.expr.clone(), + Arc::clone(&partial_sort_exec.input), + ) + .with_fetch(fetch_size), + ); + let result = + collect(Arc::new(partial_sort_exec), Arc::clone(&task_ctx)).await?; + assert_eq!( + result.iter().map(|r| r.num_rows()).collect_vec(), + expected_batch_num_rows + ); + + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + let partial_sort_result = concat_batches(&schema, &result)?; + let sort_result = collect(sort_exec, Arc::clone(&task_ctx)).await?; + assert_eq!(sort_result[0], partial_sort_result); + } + + Ok(()) + } + + #[tokio::test] + async fn test_partial_sort_no_empty_batches() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let mem_exec = prepare_partitioned_input(); + let schema = mem_exec.schema(); + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + let fetch_size = Some(250); + let partial_sort_exec = PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + Arc::clone(&mem_exec), + 1, + ) + .with_fetch(fetch_size); + + let result = collect(Arc::new(partial_sort_exec), Arc::clone(&task_ctx)).await?; + for rb in result { + assert!(rb.num_rows() > 0); + } + + Ok(()) + } + + #[tokio::test] + async fn test_sort_metadata() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let field_metadata: HashMap<String, String> = + vec![("foo".to_string(), "bar".to_string())] + .into_iter() + .collect(); + let schema_metadata: HashMap<String, String> = + vec![("baz".to_string(), "barf".to_string())] + .into_iter() + .collect(); + + let mut field = Field::new("field_name", DataType::UInt64, true); + field.set_metadata(field_metadata.clone()); + let schema = Schema::new_with_metadata(vec![field], schema_metadata.clone()); + let schema = Arc::new(schema); + + let data: ArrayRef = + Arc::new(vec![1, 1, 2].into_iter().map(Some).collect::<UInt64Array>()); + + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![data])?; + let input = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + + let partial_sort_exec = Arc::new(PartialSortExec::new( + [PhysicalSortExpr { + expr: col("field_name", &schema)?, + options: SortOptions::default(), + }] + .into(), + input, + 1, + )); + + let result: Vec<RecordBatch> = collect(partial_sort_exec, task_ctx).await?; + let expected_batch = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new( + vec![1, 1].into_iter().map(Some).collect::<UInt64Array>(), + )], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new( + vec![2].into_iter().map(Some).collect::<UInt64Array>(), + )], + )?, + ]; + + // Data is correct + assert_eq!(&expected_batch, &result); + + // explicitly ensure the metadata is present + assert_eq!(result[0].schema().fields()[0].metadata(), &field_metadata); + assert_eq!(result[0].schema().metadata(), &schema_metadata); + + Ok(()) + } + + #[tokio::test] + async fn test_lex_sort_by_float() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float32, true), + Field::new("b", DataType::Float64, true), + Field::new("c", DataType::Float64, true), + ])); + let option_asc = SortOptions { + descending: false, + nulls_first: true, + }; + let option_desc = SortOptions { + descending: true, + nulls_first: true, + }; + + // define data. + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Float32Array::from(vec![ + Some(1.0_f32), + Some(1.0_f32), + Some(1.0_f32), + Some(2.0_f32), + Some(2.0_f32), + Some(3.0_f32), + Some(3.0_f32), + Some(3.0_f32), + ])), + Arc::new(Float64Array::from(vec![ + Some(20.0_f64), + Some(20.0_f64), + Some(40.0_f64), + Some(40.0_f64), + Some(f64::NAN), + None, + None, + Some(f64::NAN), + ])), + Arc::new(Float64Array::from(vec![ + Some(10.0_f64), + Some(20.0_f64), + Some(10.0_f64), + Some(100.0_f64), + Some(f64::NAN), + Some(100.0_f64), + None, + Some(f64::NAN), + ])), + ], + )?; + + let partial_sort_exec = Arc::new(PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_desc, + }, + ] + .into(), + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)?, + 2, + )); + + assert_eq!( + DataType::Float32, + *partial_sort_exec.schema().field(0).data_type() + ); + assert_eq!( + DataType::Float64, + *partial_sort_exec.schema().field(1).data_type() + ); + assert_eq!( + DataType::Float64, + *partial_sort_exec.schema().field(2).data_type() + ); + + let result: Vec<RecordBatch> = collect( + Arc::clone(&partial_sort_exec) as Arc<dyn ExecutionPlan>, + task_ctx, + ) + .await?; + assert_snapshot!(batches_to_string(&result), @r" + +-----+------+-------+ + | a | b | c | + +-----+------+-------+ + | 1.0 | 20.0 | 20.0 | + | 1.0 | 20.0 | 10.0 | + | 1.0 | 40.0 | 10.0 | + | 2.0 | 40.0 | 100.0 | + | 2.0 | NaN | NaN | + | 3.0 | | | + | 3.0 | | 100.0 | + | 3.0 | NaN | NaN | + +-----+------+-------+ + "); + assert_eq!(result.len(), 2); + let metrics = partial_sort_exec.metrics().unwrap(); + assert!(metrics.elapsed_compute().unwrap() > 0); + assert_eq!(metrics.output_rows().unwrap(), 8); + + let columns = result[0].columns(); + + assert_eq!(DataType::Float32, *columns[0].data_type()); + assert_eq!(DataType::Float64, *columns[1].data_type()); + assert_eq!(DataType::Float64, *columns[2].data_type()); + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float32, true), + Field::new("b", DataType::Float32, true), + ])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let sort_exec = Arc::new(PartialSortExec::new( + [PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions::default(), + }] + .into(), + blocking_exec, + 1, + )); + + let fut = collect(sort_exec, Arc::clone(&task_ctx)); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partial_sort_with_homogeneous_batches() -> Result<()> { + // Test case for the bug where batches with homogeneous sort keys + // (e.g., [1,1,1], [2,2,2]) would not be properly detected as having + // slice points between batches. + let task_ctx = Arc::new(TaskContext::default()); + + // Create batches where each batch has homogeneous values for sort keys + let batch1 = test::build_table_i32( + ("a", &vec![1; 3]), + ("b", &vec![1; 3]), + ("c", &vec![3, 2, 1]), + ); + let batch2 = test::build_table_i32( + ("a", &vec![2; 3]), + ("b", &vec![2; 3]), + ("c", &vec![4, 6, 4]), + ); + let batch3 = test::build_table_i32( + ("a", &vec![3; 3]), + ("b", &vec![3; 3]), + ("c", &vec![9, 7, 8]), + ); + + let schema = batch1.schema(); + let mem_exec = TestMemoryExec::try_new_exec( + &[vec![batch1, batch2, batch3]], + Arc::clone(&schema), + None, + )?; + + let option_asc = SortOptions { + descending: false, + nulls_first: false, + }; + + // Partial sort with common prefix of 2 (sorting by a, b, c) + let partial_sort_exec = Arc::new(PartialSortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: option_asc, + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: option_asc, + }, + ] + .into(), + mem_exec, + 2, + )); + + let result = collect(partial_sort_exec, Arc::clone(&task_ctx)).await?; + + assert_eq!(result.len(), 3,); + + allow_duplicates! { + assert_snapshot!(batches_to_string(&result), @r" + +---+---+---+ + | a | b | c | + +---+---+---+ + | 1 | 1 | 1 | + | 1 | 1 | 2 | + | 1 | 1 | 3 | + | 2 | 2 | 4 | + | 2 | 2 | 4 | + | 2 | 2 | 6 | + | 3 | 3 | 7 | + | 3 | 3 | 8 | + | 3 | 3 | 9 | + +---+---+---+ + "); + } + + assert_eq!(task_ctx.runtime_env().memory_pool.reserved(), 0,); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/partitioned_topk.rs b/native/vendor/datafusion-physical-plan/src/sorts/partitioned_topk.rs new file mode 100644 index 0000000..41ccfab --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/partitioned_topk.rs
@@ -0,0 +1,530 @@ +// 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. + +//! [`PartitionedTopKExec`]: Top-K per partition operator +//! +//! For queries like: +//! ```sql +//! SELECT *, ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn +//! FROM t WHERE rn <= N +//! ``` +//! +//! Instead of sorting the entire dataset, this operator delegates to a +//! per-partition heap-of-K implementation (one variant for `ROW_NUMBER` +//! and a sibling variant for `RANK`), both of which maintain one heap per +//! distinct partition key while sharing a single [`arrow::row::RowConverter`], +//! [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation), +//! and metrics set across all partitions, and emit only the top-K rows +//! per partition in sorted order `(partition_keys, order_keys)`. + +use std::fmt::{self, Formatter}; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::row::SortField; +use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_execution::TaskContext; +use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::StreamExt; +use futures::TryStreamExt; + +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::metrics::ExecutionPlanMetricsSet; +use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields}; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, SendableRecordBatchStream, stream::RecordBatchStreamAdapter, +}; + +/// Which window function `PartitionedTopKExec` is optimizing. +/// +/// Different ranking functions have different per-partition retention rules: +/// - [`RowNumber`](Self::RowNumber): exactly K rows per partition. +/// - [`Rank`](Self::Rank): K rows plus any rows tied at the boundary +/// ORDER BY value (RANK semantics — `WHERE rk <= K` may keep more +/// than K rows when ties straddle the boundary). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WindowFnKind { + /// `ROW_NUMBER()` — keep exactly K rows per partition. + RowNumber, + /// `RANK()` — keep K rows plus any rows tied at the boundary. + Rank, +} + +/// Per-partition Top-K operator for window function queries. +/// +/// # Background +/// +/// "Top K per partition" is a common analytics pattern used for queries such as +/// "find the top 3 products by revenue for each store". The (simplified) SQL +/// for such a query might be: +/// +/// ```sql +/// SELECT * FROM ( +/// SELECT *, ROW_NUMBER() OVER (PARTITION BY store ORDER BY revenue DESC) as rn +/// FROM sales +/// ) WHERE rn <= 3; +/// ``` +/// +/// The unoptimized physical plan would be: +/// +/// ```text +/// FilterExec: rn <= 3 +/// BoundedWindowAggExec: ROW_NUMBER() PARTITION BY [store] ORDER BY [revenue DESC] +/// SortExec: expr=[store ASC, revenue DESC] +/// DataSourceExec +/// ``` +/// +/// This plan sorts the **entire** dataset (O(N log N)), computes `ROW_NUMBER` +/// for **all** rows, and then filters to keep only the top K per partition. +/// With 10M rows, 1K partitions, and K=3, it sorts all 10M rows but only +/// keeps 3K. +/// +/// # Optimization +/// +/// `PartitionedTopKExec` replaces the `SortExec` and the `FilterExec` is +/// removed. The optimized plan becomes: +/// +/// ```text +/// BoundedWindowAggExec: ROW_NUMBER() PARTITION BY [store] ORDER BY [revenue DESC] +/// PartitionedTopKExec: fetch=3, partition=[store], order=[revenue DESC] +/// DataSourceExec +/// ``` +/// +/// Instead of sorting the entire dataset, this operator reads unsorted input +/// and delegates to a per-partition heap-of-K implementation (`PartitionedTopK` +/// for `ROW_NUMBER` and `PartitionedTopKRank` for `RANK`), each maintaining +/// one heap per distinct partition key while sharing a single +/// [`arrow::row::RowConverter`] / +/// [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation) +/// across all partitions, and emits only the top-K rows per partition in +/// sorted order `(partition_keys, order_keys)`. +/// +/// Cost: O(N log K) time instead of O(N log N), and O(K × P × row_size) +/// memory where K = fetch, P = number of distinct partitions. +/// ## Why maintaining partition key order in output +/// Window functions do not require partition keys to be globally sorted, and +/// enforcing such ordering in the output can introduce unnecessary overhead. +/// However, the physical optimizer framework currently cannot express an +/// ordering that is only grouped by some keys while ordered by others. For +/// example: +/// +/// +/// # Example +/// +/// For the query above with `fetch=3` and input: +/// +/// ```text +/// store | revenue +/// ------|-------- +/// A | 100 +/// B | 50 +/// A | 200 +/// B | 150 +/// A | 300 +/// A | 400 +/// ``` +/// +/// The operator maintains two heaps: +/// - **store=A**: keeps top-3 by revenue DESC → {400, 300, 200}, evicts 100 +/// - **store=B**: keeps top-3 by revenue DESC → {150, 50} (only 2 rows) +/// +/// Output (sorted by store ASC, revenue DESC): +/// +/// ```text +/// store | revenue +/// ------|-------- +/// A | 400 +/// A | 300 +/// A | 200 +/// B | 150 +/// B | 50 +/// ``` +/// +/// This is then passed to `BoundedWindowAggExec` which assigns +/// `ROW_NUMBER` 1, 2, 3 to each partition — all of which satisfy `rn <= 3`. +/// +/// # Limitations +/// +/// - Only activated when the window function is `ROW_NUMBER` or `RANK` with +/// a `PARTITION BY` clause. `RANK` additionally requires a non-empty +/// `ORDER BY` (with an empty `ORDER BY`, every row ties at rank 1 and the +/// heap-of-K rewrite doesn't apply). Global top-K (no `PARTITION BY`) is +/// already handled efficiently by `SortExec` with `fetch`. +/// - For very high cardinality partition keys (millions of distinct values), +/// both memory usage and runtime overhead can become significant. In such +/// cases, the sort-based plan is more robust. Therefore, this optimization +/// is currently controlled by a configuration flag. +#[derive(Debug, Clone)] +pub struct PartitionedTopKExec { + /// Input execution plan (reads unsorted data) + input: Arc<dyn ExecutionPlan>, + /// Full sort expressions: `[partition_keys..., order_keys...]`. + /// + /// For `PARTITION BY store ORDER BY revenue DESC` with sort + /// `[store ASC, revenue DESC]`, the first `partition_prefix_len` + /// expressions are the partition keys (`[store ASC]`) and the + /// remaining are the order-by keys (`[revenue DESC]`). + expr: LexOrdering, + /// Number of leading expressions in `expr` that define the partition + /// key. For example, `PARTITION BY a, b` → `partition_prefix_len = 2`. + partition_prefix_len: usize, + /// Maximum number of rows to keep per partition (the K in "top-K"). + /// Derived from the filter predicate: `rn <= 3` → `fetch = 3`, + /// `rn < 3` → `fetch = 2`. + fetch: usize, + /// Which window function this operator is optimizing. Selects the + /// per-partition retention policy (see [`WindowFnKind`]). + fn_kind: WindowFnKind, + /// Execution metrics + metrics_set: ExecutionPlanMetricsSet, + /// Cached plan properties (output ordering, partitioning, etc.) + cache: Arc<PlanProperties>, +} + +impl PartitionedTopKExec { + /// Create a new `PartitionedTopKExec`. + /// + /// # Arguments + /// + /// * `input` - The child execution plan providing unsorted input rows. + /// * `expr` - Full sort ordering `[partition_keys..., order_keys...]`. + /// For `PARTITION BY pk ORDER BY val ASC`, this would be `[pk ASC, val ASC]`. + /// * `partition_prefix_len` - Number of leading expressions in `expr` + /// that form the partition key. Must be >= 1. + /// * `fetch` - Maximum rows to retain per partition (the K in "top-K"). + /// * `fn_kind` - Which ranking window function this operator optimizes + /// ([`WindowFnKind::RowNumber`] or [`WindowFnKind::Rank`]). + /// + /// # Example + /// + /// ```text + /// // For: ROW_NUMBER() OVER (PARTITION BY store ORDER BY revenue DESC) ... WHERE rn <= 5 + /// PartitionedTopKExec::try_new( + /// data_source, + /// LexOrdering([store ASC, revenue DESC]), + /// 1, // partition_prefix_len: 1 partition column (store) + /// 5, // fetch: keep top 5 per partition + /// WindowFnKind::RowNumber, + /// ) + /// ``` + pub fn try_new( + input: Arc<dyn ExecutionPlan>, + expr: LexOrdering, + partition_prefix_len: usize, + fetch: usize, + fn_kind: WindowFnKind, + ) -> Result<Self> { + let cache = Self::compute_properties(&input, expr.clone())?; + Ok(Self { + input, + expr, + partition_prefix_len, + fetch, + fn_kind, + metrics_set: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + }) + } + + /// Returns the child execution plan. + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Returns the full sort ordering `[partition_keys..., order_keys...]`. + pub fn expr(&self) -> &LexOrdering { + &self.expr + } + + /// Returns the number of leading expressions in [`Self::expr`] that + /// define the partition key. + pub fn partition_prefix_len(&self) -> usize { + self.partition_prefix_len + } + + /// Returns the maximum number of rows retained per partition. + pub fn fetch(&self) -> usize { + self.fetch + } + + /// Returns which window function this operator is optimizing. + pub fn fn_kind(&self) -> WindowFnKind { + self.fn_kind + } + + /// Compute [`PlanProperties`] for this operator. + /// + /// The output is sorted by `sort_exprs` (partition keys then order keys), + /// uses the same partitioning as the input, emits all output at once + /// (`EmissionType::Final`), and is bounded. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + sort_exprs: LexOrdering, + ) -> Result<PlanProperties> { + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.reorder(sort_exprs)?; + + Ok(PlanProperties::new( + eq_properties, + input.output_partitioning().clone(), + EmissionType::Final, + Boundedness::Bounded, + )) + } +} + +impl DisplayAs for PartitionedTopKExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + let fn_label = match self.fn_kind { + WindowFnKind::RowNumber => "row_number", + WindowFnKind::Rank => "rank", + }; + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len] + .iter() + .map(|e| format!("{}", e.expr)) + .collect(); + let order_exprs: Vec<String> = self.expr[self.partition_prefix_len..] + .iter() + .map(|e| format!("{e}")) + .collect(); + write!( + f, + "PartitionedTopKExec: fn={}, fetch={}, partition=[{}], order=[{}]", + fn_label, + self.fetch, + partition_exprs.join(", "), + order_exprs.join(", "), + ) + } + DisplayFormatType::TreeRender => { + let partition_exprs: Vec<String> = self.expr[..self.partition_prefix_len] + .iter() + .map(|e| format!("{}", e.expr)) + .collect(); + let order_exprs: Vec<String> = self.expr[self.partition_prefix_len..] + .iter() + .map(|e| format!("{e}")) + .collect(); + writeln!(f, "fn={fn_label}")?; + writeln!(f, "fetch={}", self.fetch)?; + writeln!(f, "partition=[{}]", partition_exprs.join(", "))?; + writeln!(f, "order=[{}]", order_exprs.join(", ")) + } + } + } +} + +impl ExecutionPlan for PartitionedTopKExec { + fn name(&self) -> &'static str { + "PartitionedTopKExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + let partition_exprs: Vec<Arc<dyn PhysicalExpr>> = self.expr + [..self.partition_prefix_len] + .iter() + .map(|e| Arc::clone(&e.expr)) + .collect(); + crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + partition_exprs, + )]) + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![false] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + assert_eq!(children.len(), 1); + Ok(Arc::new(PartitionedTopKExec::try_new( + Arc::clone(&children[0]), + self.expr.clone(), + self.partition_prefix_len, + self.fetch, + self.fn_kind, + )?)) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, Arc::clone(&context))?; + let schema = input.schema(); + + let partition_sort_fields = + build_sort_fields(&self.expr[..self.partition_prefix_len], &schema)?; + + let partition_exprs: Vec<Arc<dyn PhysicalExpr>> = self.expr + [..self.partition_prefix_len] + .iter() + .map(|e| Arc::clone(&e.expr)) + .collect(); + let order_expr: LexOrdering = + LexOrdering::new(self.expr[self.partition_prefix_len..].iter().cloned()) + .expect("PartitionedTopKExec requires at least one order-by expression"); + let fetch = self.fetch; + let fn_kind = self.fn_kind; + let batch_size = context.session_config().batch_size(); + let runtime = Arc::clone(&context.runtime_env()); + let metrics_set = self.metrics_set.clone(); + + let stream = futures::stream::once(async move { + do_partitioned_topk( + partition, + input, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + fetch, + fn_kind, + batch_size, + runtime, + metrics_set, + ) + .await + }) + .try_flatten(); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.input.schema(), + stream, + ))) + } +} + +/// Read all input, feed each batch into a per-partition top-K state +/// (either [`PartitionedTopK`] for `ROW_NUMBER` or +/// [`PartitionedTopKRank`] for `RANK`), then emit results ordered by +/// `(partition_keys, order_keys)`. +/// +/// # Phases +/// +/// 1. **Accumulation** — forward each input `RecordBatch` to the +/// per-partition state's `insert_batch`. The `RowConverter` for +/// ORDER BY columns, the operator's `MemoryReservation`, and the +/// `TopKMetrics` are shared across all distinct partition keys for +/// this operator instance. +/// +/// 2. **Emission** — `emit` drains all per-partition heaps in sorted +/// partition-key order, returning a coalesced batch stream. For +/// `RANK`, boundary-tied rows are materialized and emitted after +/// each partition's heap rows. +/// +/// # Cost +/// +/// - Time: O(N log K) where N = total rows, K = fetch +/// - Memory: O(K × P × row_size) where P = number of distinct partitions +/// plus, for RANK, the boundary ties' rows +#[expect(clippy::too_many_arguments)] +async fn do_partitioned_topk( + partition_id: usize, + mut input: SendableRecordBatchStream, + schema: SchemaRef, + partition_exprs: Vec<Arc<dyn PhysicalExpr>>, + partition_sort_fields: Vec<SortField>, + order_expr: LexOrdering, + fetch: usize, + fn_kind: WindowFnKind, + batch_size: usize, + runtime: Arc<RuntimeEnv>, + metrics_set: ExecutionPlanMetricsSet, +) -> Result<SendableRecordBatchStream> { + match fn_kind { + WindowFnKind::RowNumber => { + let mut state = PartitionedTopK::try_new( + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + fetch, + batch_size, + &runtime, + &metrics_set, + )?; + while let Some(batch) = input.next().await { + state.insert_batch(&batch?)?; + } + drop(input); + state.emit() + } + WindowFnKind::Rank => { + let mut state = PartitionedTopKRank::try_new( + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + fetch, + batch_size, + &runtime, + &metrics_set, + )?; + while let Some(batch) = input.next().await { + state.insert_batch(&batch?)?; + } + drop(input); + state.emit() + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs new file mode 100644 index 0000000..6c782f5 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs
@@ -0,0 +1,3627 @@ +// 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 that deals with an arbitrary size of the input. +//! It will do in-memory sorting if it has enough memory budget +//! but spills to disk if needed. + +use std::fmt; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use parking_lot::RwLock; + +use crate::common::spawn_buffered; +use crate::execution_plan::{ + Boundedness, CardinalityEffect, EmissionType, has_same_children_properties, + replace_children_if_necessary, +}; +use crate::expressions::PhysicalSortExpr; +use crate::filter::FilterExec; +use crate::filter_pushdown::{ + ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, PushedDown, +}; +use crate::limit::LimitStream; +use crate::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, SpillMetrics, +}; +use crate::projection::{ProjectionExec, make_with_child, update_ordering}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::get_record_batch_memory_size; +use crate::spill::in_progress_spill_file::InProgressSpillFile; +use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::ReservationStream; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; +use crate::topk::TopK; +use crate::topk::TopKDynamicFilters; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, + EmptyRecordBatchStream, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, +}; + +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays}; +use arrow::datatypes::SchemaRef; +use datafusion_common::config::SpillCompression; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + DataFusionError, Result, assert_or_internal_err, internal_datafusion_err, + unwrap_or_internal_err, +}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_physical_expr::LexOrdering; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; + +use futures::{StreamExt, TryStreamExt}; +use log::{debug, trace}; + +struct ExternalSorterMetrics { + /// metrics + baseline: BaselineMetrics, + + spill_metrics: SpillMetrics, +} + +impl ExternalSorterMetrics { + fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + spill_metrics: SpillMetrics::new(metrics, partition), + } + } +} + +/// Sorts an arbitrary sized, unsorted, stream of [`RecordBatch`]es to +/// a total order. Depending on the input size and memory manager +/// configuration, writes intermediate results to disk ("spills") +/// using Arrow IPC format. +/// +/// # Algorithm +/// +/// 1. get a non-empty new batch from input +/// +/// 2. check with the memory manager there is sufficient space to +/// buffer the batch in memory. +/// +/// 2.1 if memory is sufficient, buffer batch in memory, go to 1. +/// +/// 2.2 if no more memory is available, sort all buffered batches and +/// spill to file. buffer the next batch in memory, go to 1. +/// +/// 3. when input is exhausted, merge all in memory batches and spills +/// to get a total order. +/// +/// # When data fits in available memory +/// +/// If there is sufficient memory, data is sorted in memory to produce the output +/// +/// ```text +/// ┌─────┐ +/// │ 2 │ +/// │ 3 │ +/// │ 1 │─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ +/// │ 4 │ +/// │ 2 │ │ +/// └─────┘ ▼ +/// ┌─────┐ +/// │ 1 │ In memory +/// │ 4 │─ ─ ─ ─ ─ ─▶ sort/merge ─ ─ ─ ─ ─▶ total sorted output +/// │ 1 │ +/// └─────┘ ▲ +/// ... │ +/// +/// ┌─────┐ │ +/// │ 4 │ +/// │ 3 │─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +/// └─────┘ +/// +/// in_mem_batches +/// ``` +/// +/// # When data does not fit in available memory +/// +/// When memory is exhausted, data is first sorted and written to one +/// or more spill files on disk: +/// +/// ```text +/// ┌─────┐ .─────────────────. +/// │ 2 │ ( ) +/// │ 3 │ │`─────────────────'│ +/// │ 1 │─ ─ ─ ─ ─ ─ ─ │ ┌────┐ │ +/// │ 4 │ │ │ │ 1 │░ │ +/// │ 2 │ │ │... │░ │ +/// └─────┘ ▼ │ │ 4 │░ ┌ ─ ─ │ +/// ┌─────┐ │ └────┘░ 1 │░ │ +/// │ 1 │ In memory │ ░░░░░░ │ ░░ │ +/// │ 4 │─ ─ ▶ sort/merge ─ ─ ─ ─ ┼ ─ ─ ─ ─ ─▶ ... │░ │ +/// │ 1 │ and write to file │ │ ░░ │ +/// └─────┘ │ 4 │░ │ +/// ... ▲ │ └░─░─░░ │ +/// │ │ ░░░░░░ │ +/// ┌─────┐ │.─────────────────.│ +/// │ 4 │ │ ( ) +/// │ 3 │─ ─ ─ ─ ─ ─ ─ `─────────────────' +/// └─────┘ +/// +/// in_mem_batches spills +/// (file on disk in Arrow +/// IPC format) +/// ``` +/// +/// Once the input is completely read, the spill files are read and +/// merged with any in memory batches to produce a single total sorted +/// output: +/// +/// ```text +/// .─────────────────. +/// ( ) +/// │`─────────────────'│ +/// │ ┌────┐ │ +/// │ │ 1 │░ │ +/// │ │... │─ ─ ─ ─ ─ ─│─ ─ ─ ─ ─ ─ +/// │ │ 4 │░ ┌────┐ │ │ +/// │ └────┘░ │ 1 │░ │ ▼ +/// │ ░░░░░░ │ │░ │ +/// │ │... │─ ─│─ ─ ─ ▶ merge ─ ─ ─▶ total sorted output +/// │ │ │░ │ +/// │ │ 4 │░ │ ▲ +/// │ └────┘░ │ │ +/// │ ░░░░░░ │ +/// │.─────────────────.│ │ +/// ( ) +/// `─────────────────' │ +/// spills +/// │ +/// +/// │ +/// +/// ┌─────┐ │ +/// │ 1 │ +/// │ 4 │─ ─ ─ ─ │ +/// └─────┘ │ +/// ... In memory +/// └ ─ ─ ─▶ sort/merge +/// ┌─────┐ +/// │ 4 │ ▲ +/// │ 3 │─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +/// └─────┘ +/// +/// in_mem_batches +/// ``` +struct ExternalSorter { + // ======================================================================== + // PROPERTIES: + // Fields that define the sorter's configuration and remain constant + // ======================================================================== + /// Schema of the output (and the input) + schema: SchemaRef, + /// Sort expressions + expr: LexOrdering, + /// The target number of rows for output batches + batch_size: usize, + /// If the in size of buffered memory batches is below this size, + /// the data will be concatenated and sorted in place rather than + /// sort/merged. + sort_in_place_threshold_bytes: usize, + + // ======================================================================== + // STATE BUFFERS: + // Fields that hold intermediate data during sorting + // ======================================================================== + /// Unsorted input batches stored in the memory buffer + in_mem_batches: Vec<RecordBatch>, + + /// During external sorting, in-memory intermediate data will be appended to + /// this file incrementally. Once finished, this file will be moved to [`Self::finished_spill_files`]. + /// + /// this is a tuple of: + /// 1. `InProgressSpillFile` - the file that is being written to + /// 2. `max_record_batch_memory` - the maximum memory usage of a single batch in this spill file. + in_progress_spill_file: Option<(InProgressSpillFile, usize)>, + /// If data has previously been spilled, the locations of the spill files (in + /// Arrow IPC format) + /// Within the same spill file, the data might be chunked into multiple batches, + /// and ordered by sort keys. + finished_spill_files: Vec<SortedSpillFile>, + + // ======================================================================== + // EXECUTION RESOURCES: + // Fields related to managing execution resources and monitoring performance. + // ======================================================================== + /// Runtime metrics + metrics: ExternalSorterMetrics, + /// A handle to the runtime to get spill files + runtime: Arc<RuntimeEnv>, + /// Reservation for in_mem_batches + reservation: MemoryReservation, + spill_manager: SpillManager, + + /// Reservation for the merging of in-memory batches. If the sort + /// might spill, `sort_spill_reservation_bytes` will be + /// pre-reserved to ensure there is some space for this sort/merge. + merge_reservation: MemoryReservation, + /// How much memory to reserve for performing in-memory sort/merges + /// prior to spilling. + sort_spill_reservation_bytes: usize, +} + +impl ExternalSorter { + // TODO: make a builder or some other nicer API to avoid the + // clippy warning + #[expect(clippy::too_many_arguments)] + pub fn new( + partition_id: usize, + schema: SchemaRef, + expr: LexOrdering, + batch_size: usize, + sort_spill_reservation_bytes: usize, + sort_in_place_threshold_bytes: usize, + // Configured via `datafusion.execution.spill_compression`. + spill_compression: SpillCompression, + metrics: &ExecutionPlanMetricsSet, + runtime: Arc<RuntimeEnv>, + ) -> Result<Self> { + let metrics = ExternalSorterMetrics::new(metrics, partition_id); + let reservation = MemoryConsumer::new(format!("ExternalSorter[{partition_id}]")) + .with_can_spill(true) + .register(&runtime.memory_pool); + + let merge_reservation = + MemoryConsumer::new(format!("ExternalSorterMerge[{partition_id}]")) + .register(&runtime.memory_pool); + + let spill_manager = SpillManager::new( + Arc::clone(&runtime), + metrics.spill_metrics.clone(), + Arc::clone(&schema), + ) + .with_compression_type(spill_compression); + + Ok(Self { + schema, + in_mem_batches: vec![], + in_progress_spill_file: None, + finished_spill_files: vec![], + expr, + metrics, + reservation, + spill_manager, + merge_reservation, + runtime, + batch_size, + sort_spill_reservation_bytes, + sort_in_place_threshold_bytes, + }) + } + + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` + /// + /// Updates memory usage metrics, and possibly triggers spilling to disk + async fn insert_batch(&mut self, input: RecordBatch) -> Result<()> { + if input.num_rows() == 0 { + return Ok(()); + } + + self.reserve_memory_for_merge()?; + self.reserve_memory_for_batch_and_maybe_spill(&input) + .await?; + + self.in_mem_batches.push(input); + Ok(()) + } + + fn spilled_before(&self) -> bool { + !self.finished_spill_files.is_empty() + } + + /// Returns the final sorted output of all batches inserted via + /// [`Self::insert_batch`] as a stream of [`RecordBatch`]es. + /// + /// This process could either be: + /// + /// 1. An in-memory sort/merge (if the input fit in memory) + /// + /// 2. A combined streaming merge incorporating both in-memory + /// batches and data from spill files on disk. + async fn sort(&mut self) -> Result<SendableRecordBatchStream> { + if self.spilled_before() { + // Sort `in_mem_batches` and spill it first. If there are many + // `in_mem_batches` and the memory limit is almost reached, merging + // them with the spilled files at the same time might cause OOM. + if !self.in_mem_batches.is_empty() { + self.sort_and_spill_in_mem_batches().await?; + } + + // Transfer the pre-reserved merge memory to the streaming merge + // using `take()` instead of `new_empty()`. This ensures the merge + // stream starts with `sort_spill_reservation_bytes` already + // allocated, preventing starvation when concurrent sort partitions + // compete for pool memory. `take()` moves the bytes atomically + // without releasing them back to the pool, so other partitions + // cannot race to consume the freed memory. + StreamingMergeBuilder::new() + .with_sorted_spill_files(std::mem::take(&mut self.finished_spill_files)) + .with_spill_manager(self.spill_manager.clone()) + .with_schema(Arc::clone(&self.schema)) + .with_expressions(&self.expr.clone()) + .with_metrics(self.metrics.baseline.clone()) + .with_batch_size(self.batch_size) + .with_fetch(None) + .with_reservation(self.merge_reservation.take()) + .build() + } else { + // Release the memory reserved for merge back to the pool so + // there is some left when `in_mem_sort_stream` requests an + // allocation. Only needed for the non-spill path; the spill + // path transfers the reservation to the merge stream instead. + self.merge_reservation.free(); + self.in_mem_sort_stream(true, true) + } + } + + /// How much memory is buffered in this `ExternalSorter`? + fn used(&self) -> usize { + self.reservation.size() + } + + /// How much memory is reserved for the merge phase? + #[cfg(test)] + fn merge_reservation_size(&self) -> usize { + self.merge_reservation.size() + } + + /// How many bytes have been spilled to disk? + fn spilled_bytes(&self) -> usize { + self.metrics.spill_metrics.spilled_bytes.value() + } + + /// How many rows have been spilled to disk? + fn spilled_rows(&self) -> usize { + self.metrics.spill_metrics.spilled_rows.value() + } + + /// How many spill files have been created? + fn spill_count(&self) -> usize { + self.metrics.spill_metrics.spill_file_count.value() + } + + /// Appending globally sorted batches to the in-progress spill file, and clears + /// the `globally_sorted_batches` (also its memory reservation) afterwards. + fn consume_and_spill_append( + &mut self, + globally_sorted_batches: &mut Vec<RecordBatch>, + ) -> Result<()> { + if globally_sorted_batches.is_empty() { + return Ok(()); + } + + // Lazily initialize the in-progress spill file + if self.in_progress_spill_file.is_none() { + self.in_progress_spill_file = + Some((self.spill_manager.create_in_progress_file("Sorting")?, 0)); + } + + debug!("Spilling sort data of ExternalSorter to disk whilst inserting"); + + let batches_to_spill = std::mem::take(globally_sorted_batches); + self.reservation.free(); + + let (in_progress_file, max_record_batch_size) = + self.in_progress_spill_file.as_mut().ok_or_else(|| { + internal_datafusion_err!("In-progress spill file should be initialized") + })?; + + for batch in batches_to_spill { + let gc_sliced_size = in_progress_file.append_batch(&batch)?; + + *max_record_batch_size = (*max_record_batch_size).max(gc_sliced_size); + } + + assert_or_internal_err!( + globally_sorted_batches.is_empty(), + "This function consumes globally_sorted_batches, so it should be empty after taking." + ); + + Ok(()) + } + + /// Finishes the in-progress spill file and moves it to the finished spill files. + fn spill_finish(&mut self) -> Result<()> { + let (mut in_progress_file, max_record_batch_memory) = + self.in_progress_spill_file.take().ok_or_else(|| { + internal_datafusion_err!("Should be called after `spill_append`") + })?; + let spill_file = in_progress_file.finish()?; + + if let Some(spill_file) = spill_file { + self.finished_spill_files.push(SortedSpillFile { + file: spill_file, + max_record_batch_memory, + }); + } + + Ok(()) + } + + /// Sorts the in-memory batches and merges them into a single sorted run, then writes + /// the result to spill files. + async fn sort_and_spill_in_mem_batches(&mut self) -> Result<()> { + assert_or_internal_err!( + !self.in_mem_batches.is_empty(), + "in_mem_batches must not be empty when attempting to sort and spill" + ); + + // Release the memory reserved for merge back to the pool so + // there is some left when `in_mem_sort_stream` requests an + // allocation. At the end of this function, memory will be + // reserved again for the next spill. + self.merge_reservation.free(); + + let mut sorted_stream = self.in_mem_sort_stream( + false, + // No coalescing on the spill path: it raises per-run peak memory. + false, + )?; + // After `in_mem_sort_stream()` is constructed, all `in_mem_batches` is taken + // to construct a globally sorted stream. + assert_or_internal_err!( + self.in_mem_batches.is_empty(), + "in_mem_batches should be empty after constructing sorted stream" + ); + // 'global' here refers to all buffered batches when the memory limit is + // reached. This variable will buffer the sorted batches after + // sort-preserving merge and incrementally append to spill files. + let mut globally_sorted_batches: Vec<RecordBatch> = vec![]; + + while let Some(batch) = sorted_stream.next().await { + let batch = batch?; + let sorted_size = get_reserved_bytes_for_record_batch(&batch)?; + if self.reservation.try_grow(sorted_size).is_err() { + // Although the reservation is not enough, the batch is + // already in memory, so it's okay to combine it with previously + // sorted batches, and spill together. + globally_sorted_batches.push(batch); + self.consume_and_spill_append(&mut globally_sorted_batches)?; // reservation is freed in spill() + } else { + globally_sorted_batches.push(batch); + } + } + + // Drop early to free up memory reserved by the sorted stream, otherwise the + // upcoming `self.reserve_memory_for_merge()` may fail due to insufficient memory. + drop(sorted_stream); + + self.consume_and_spill_append(&mut globally_sorted_batches)?; + self.spill_finish()?; + + // Sanity check after spilling + let buffers_cleared_property = + self.in_mem_batches.is_empty() && globally_sorted_batches.is_empty(); + assert_or_internal_err!( + buffers_cleared_property, + "in_mem_batches and globally_sorted_batches should be cleared before" + ); + + // Reserve headroom for next sort/merge + self.reserve_memory_for_merge()?; + + Ok(()) + } + + /// Consumes in_mem_batches returning a sorted stream of + /// batches. This proceeds in one of two ways: + /// + /// # Small Datasets + /// + /// For "smaller" datasets, the data is first concatenated into a + /// single batch and then sorted. This is often faster than + /// sorting and then merging. + /// + /// ```text + /// ┌─────┐ + /// │ 2 │ + /// │ 3 │ + /// │ 1 │─ ─ ─ ─ ┐ ┌─────┐ + /// │ 4 │ │ 2 │ + /// │ 2 │ │ │ 3 │ + /// └─────┘ │ 1 │ sorted output + /// ┌─────┐ ▼ │ 4 │ stream + /// │ 1 │ │ 2 │ + /// │ 4 │─ ─▶ concat ─ ─ ─ ─ ▶│ 1 │─ ─ ▶ sort ─ ─ ─ ─ ─▶ + /// │ 1 │ │ 4 │ + /// └─────┘ ▲ │ 1 │ + /// ... │ │ ... │ + /// │ 4 │ + /// ┌─────┐ │ │ 3 │ + /// │ 4 │ └─────┘ + /// │ 3 │─ ─ ─ ─ ┘ + /// └─────┘ + /// in_mem_batches + /// ``` + /// + /// # Larger datasets + /// + /// For larger datasets, the batches are first sorted individually + /// and then merged together. + /// + /// ```text + /// ┌─────┐ ┌─────┐ + /// │ 2 │ │ 1 │ + /// │ 3 │ │ 2 │ + /// │ 1 │─ ─▶ sort ─ ─▶│ 2 │─ ─ ─ ─ ─ ┐ + /// │ 4 │ │ 3 │ + /// │ 2 │ │ 4 │ │ + /// └─────┘ └─────┘ sorted output + /// ┌─────┐ ┌─────┐ ▼ stream + /// │ 1 │ │ 1 │ + /// │ 4 │─ ▶ sort ─ ─ ▶│ 1 ├ ─ ─ ▶ merge ─ ─ ─ ─▶ + /// │ 1 │ │ 4 │ + /// └─────┘ └─────┘ ▲ + /// ... ... ... │ + /// + /// ┌─────┐ ┌─────┐ │ + /// │ 4 │ │ 3 │ + /// │ 3 │─ ▶ sort ─ ─ ▶│ 4 │─ ─ ─ ─ ─ ┘ + /// └─────┘ └─────┘ + /// + /// in_mem_batches + /// ``` + /// `coalesce_runs` merges buffered batches into fewer, larger sorted runs to + /// reduce merge fan-in. Disabled on the spill path to keep peak memory low. + fn in_mem_sort_stream( + &mut self, + is_output_stream: bool, + coalesce_runs: bool, + ) -> Result<SendableRecordBatchStream> { + if self.in_mem_batches.is_empty() { + let empty_stream = + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); + return Ok(self.observe_if_output(empty_stream, is_output_stream)); + } + + // The elapsed compute timer is updated when the value is dropped. + // There is no need for an explicit call to drop. + let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); + + // Please pay attention that any operation inside of `in_mem_sort_stream` will + // not perform any memory reservation. This is for avoiding the need of handling + // reservation failure and spilling in the middle of the sort/merge. The memory + // space for batches produced by the resulting stream will be reserved by the + // consumer of the stream. + + if self.in_mem_batches.len() == 1 { + let batch = self.in_mem_batches.swap_remove(0); + let reservation = self.reservation.take(); + let sorted_stream = self.sort_batch_stream(batch, reservation)?; + return Ok(self.observe_if_output(sorted_stream, is_output_stream)); + } + + // If less than sort_in_place_threshold_bytes, concatenate and sort in place + if self.reservation.size() < self.sort_in_place_threshold_bytes { + // Concatenate memory batches together and sort + let batch = concat_batches(&self.schema, &self.in_mem_batches)?; + self.in_mem_batches.clear(); + self.reservation + .try_resize(get_reserved_bytes_for_record_batch(&batch)?) + .map_err(Self::err_with_oom_context)?; + let reservation = self.reservation.take(); + let sorted_stream = self.sort_batch_stream(batch, reservation)?; + return Ok(self.observe_if_output(sorted_stream, is_output_stream)); + } + + // For single-column sorts, coalesce the buffered batches into fewer, + // larger runs to cut the merge fan-in (where the cheap per-key compare is + // dominated by per-stream cursor/merge overhead). Multi-column sorts are + // left as one run per batch: the row-format merge of many small runs + // beats sorting a few large runs with the lexicographic comparator. + let batches = std::mem::take(&mut self.in_mem_batches); + let runs = if coalesce_runs && self.expr.len() == 1 { + self.coalesce_in_mem_batches_into_runs(batches)? + } else { + batches + }; + + let streams = runs + .into_iter() + .map(|batch| { + let reservation = self + .reservation + .split(get_reserved_bytes_for_record_batch(&batch)?); + let input = self.sort_batch_stream(batch, reservation)?; + Ok(spawn_buffered(input, 1)) + }) + .collect::<Result<_>>()?; + + StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(Arc::clone(&self.schema)) + .with_expressions(&self.expr.clone()) + .with_metrics(if is_output_stream { + self.metrics.baseline.clone() + } else { + self.metrics.baseline.intermediate() + }) + .with_batch_size(self.batch_size) + .with_fetch(None) + .with_reservation(self.merge_reservation.new_empty()) + .build() + } + + /// Concatenates `batches` into fewer, larger runs, each bounded by + /// `sort_in_place_threshold_bytes`, to reduce merge fan-in. `self.reservation` + /// is resized to the coalesced footprint so the caller's per-run splits stay + /// exact. + fn coalesce_in_mem_batches_into_runs( + &mut self, + batches: Vec<RecordBatch>, + ) -> Result<Vec<RecordBatch>> { + let target = self.sort_in_place_threshold_bytes.max(1); + let mut runs: Vec<RecordBatch> = Vec::new(); + let mut group: Vec<RecordBatch> = Vec::new(); + let mut group_bytes = 0usize; + + // Flush a group into a run, skipping the copy for a single-batch group. + let flush = |group: &mut Vec<RecordBatch>, + runs: &mut Vec<RecordBatch>, + schema: &SchemaRef| + -> Result<()> { + match group.len() { + 0 => {} + 1 => runs.push(group.pop().unwrap()), + _ => { + runs.push(concat_batches(schema, group.iter())?); + group.clear(); + } + } + Ok(()) + }; + + for batch in batches { + let bytes = get_reserved_bytes_for_record_batch(&batch)?; + if !group.is_empty() && group_bytes.saturating_add(bytes) > target { + flush(&mut group, &mut runs, &self.schema)?; + group_bytes = 0; + } + group_bytes += bytes; + group.push(batch); + } + flush(&mut group, &mut runs, &self.schema)?; + + // Realign the reservation: concatenation may shift the footprint slightly. + let total: usize = runs + .iter() + .map(get_reserved_bytes_for_record_batch) + .sum::<Result<usize>>()?; + self.reservation + .try_resize(total) + .map_err(Self::err_with_oom_context)?; + + Ok(runs) + } + + /// Sorts a single `RecordBatch` into a single stream. + /// + /// This may output multiple batches depending on the size of the + /// sorted data and the target batch size. + /// For single-batch output cases, `reservation` will be freed immediately after sorting, + /// as the batch will be output and is expected to be reserved by the consumer of the stream. + /// For multi-batch output cases, `reservation` will be grown to match the actual + /// size of sorted output, and as each batch is output, its memory will be freed from the reservation. + /// (This leads to the same behaviour, as futures are only evaluated when polled by the consumer.) + fn sort_batch_stream( + &self, + batch: RecordBatch, + reservation: MemoryReservation, + ) -> Result<SendableRecordBatchStream> { + assert_eq!( + get_reserved_bytes_for_record_batch(&batch)?, + reservation.size() + ); + + let schema = batch.schema(); + let expressions = self.expr.clone(); + let batch_size = self.batch_size; + + let stream = futures::stream::once(async move { + let schema = batch.schema(); + + // Sort the batch immediately and get all output batches + let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; + + // Resize the reservation to match the actual sorted output size. + // Using try_resize avoids a release-then-reacquire cycle, which + // matters for MemoryPool implementations where grow/shrink have + // non-trivial cost (e.g. JNI calls in Comet). + let total_sorted_size: usize = sorted_batches + .iter() + .map(get_record_batch_memory_size) + .sum(); + reservation + .try_resize(total_sorted_size) + .map_err(Self::err_with_oom_context)?; + + // Wrap in ReservationStream to hold the reservation + Result::<_, DataFusionError>::Ok(Box::pin(ReservationStream::new( + Arc::clone(&schema), + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(sorted_batches.into_iter().map(Ok)), + )), + reservation, + )) as SendableRecordBatchStream) + }) + .try_flatten(); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + /// If this sort may spill, pre-allocates + /// `sort_spill_reservation_bytes` of memory to guarantee memory + /// left for the in memory sort/merge. + fn reserve_memory_for_merge(&mut self) -> Result<()> { + // Reserve headroom for next merge sort + if self.runtime.disk_manager.tmp_files_enabled() { + let size = self.sort_spill_reservation_bytes; + if self.merge_reservation.size() != size { + self.merge_reservation + .try_resize(size) + .map_err(Self::err_with_oom_context)?; + } + } + + Ok(()) + } + + /// Reserves memory to be able to accommodate the given batch. + /// If memory is scarce, tries to spill current in-memory batches to disk first. + async fn reserve_memory_for_batch_and_maybe_spill( + &mut self, + input: &RecordBatch, + ) -> Result<()> { + let size = get_reserved_bytes_for_record_batch(input)?; + + match self.reservation.try_grow(size) { + Ok(_) => Ok(()), + Err(e) => { + if self.in_mem_batches.is_empty() { + return Err(Self::err_with_oom_context(e)); + } + + // Spill and try again. + self.sort_and_spill_in_mem_batches().await?; + self.reservation + .try_grow(size) + .map_err(Self::err_with_oom_context) + } + } + } + + /// Wraps the error with a context message suggesting settings to tweak. + /// This is meant to be used with DataFusionError::ResourcesExhausted only. + fn err_with_oom_context(e: DataFusionError) -> DataFusionError { + match e { + DataFusionError::ResourcesExhausted(_) => e.context( + "Not enough memory to continue external sort. \ + Consider increasing the memory limit config: 'datafusion.runtime.memory_limit', \ + or decreasing the config: 'datafusion.execution.sort_spill_reservation_bytes'." + ), + // This is not an OOM error, so just return it as is. + _ => e, + } + } + + fn observe_if_output( + &self, + mut stream: SendableRecordBatchStream, + wrap: bool, + ) -> SendableRecordBatchStream { + if wrap { + stream = Box::pin(ObservedStream::new( + stream, + self.metrics.baseline.clone(), + None, + )) + } + + stream + } +} + +/// Estimate how much memory is needed to sort a `RecordBatch`. +/// +/// This is used to pre-reserve memory for the sort/merge. The sort/merge process involves +/// creating sorted copies of sorted columns in record batches for speeding up comparison +/// in sorting and merging. The sorted copies are in either row format or array format. +/// Please refer to cursor.rs and stream.rs for more details. No matter what format the +/// sorted copies are, they will use more memory than the original record batch. +/// +/// This can basically be calculated as the sum of the actual space it takes in +/// memory (which would be larger for a sliced batch), and the size of the actual data. +pub(crate) fn get_reserved_bytes_for_record_batch_size( + record_batch_size: usize, + sliced_size: usize, +) -> usize { + // Even 2x may not be enough for some cases, but it's a good enough estimation as a baseline. + // If 2x is not enough, user can set a larger value for `sort_spill_reservation_bytes` + // to compensate for the extra memory needed. + record_batch_size + sliced_size +} + +/// Estimate how much memory is needed to sort a `RecordBatch`. +/// This will just call `get_reserved_bytes_for_record_batch_size` with the +/// memory size of the record batch and its sliced size. +pub(crate) fn get_reserved_bytes_for_record_batch(batch: &RecordBatch) -> Result<usize> { + batch.get_sliced_size().map(|sliced_size| { + get_reserved_bytes_for_record_batch_size( + get_record_batch_memory_size(batch), + sliced_size, + ) + }) +} + +impl Debug for ExternalSorter { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("ExternalSorter") + .field("memory_used", &self.used()) + .field("spilled_bytes", &self.spilled_bytes()) + .field("spilled_rows", &self.spilled_rows()) + .field("spill_count", &self.spill_count()) + .finish() + } +} + +pub fn sort_batch( + batch: &RecordBatch, + expressions: &LexOrdering, + fetch: Option<usize>, +) -> Result<RecordBatch> { + let sort_columns = expressions + .iter() + .map(|expr| expr.evaluate_to_sort_column(batch)) + .collect::<Result<Vec<_>>>()?; + + let indices = lexsort_to_indices(&sort_columns, fetch)?; + let columns = take_arrays(batch.columns(), &indices, None)?; + + let options = RecordBatchOptions::new().with_row_count(Some(indices.len())); + Ok(RecordBatch::try_new_with_options( + batch.schema(), + columns, + &options, + )?) +} + +/// Sort a batch and return the result as multiple batches of size `batch_size`. +/// This is useful when you want to avoid creating one large sorted batch in memory, +/// and instead want to process the sorted data in smaller chunks. +pub fn sort_batch_chunked( + batch: &RecordBatch, + expressions: &LexOrdering, + batch_size: usize, +) -> Result<Vec<RecordBatch>> { + IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size).collect() +} + +/// Sort execution plan. +/// +/// Support sorting datasets that are larger than the memory allotted +/// by the memory manager, by spilling to disk. +#[derive(Debug, Clone)] +pub struct SortExec { + /// Input schema + pub(crate) input: Arc<dyn ExecutionPlan>, + /// Sort expressions + expr: LexOrdering, + /// Containing all metrics set created during sort + metrics_set: ExecutionPlanMetricsSet, + /// Preserve partitions of input plan. If false, the input partitions + /// will be sorted and merged into a single output partition. + preserve_partitioning: bool, + /// Fetch highest/lowest n results + fetch: Option<usize>, + /// Normalized common sort prefix between the input and the sort expressions (only used with fetch) + common_sort_prefix: Vec<PhysicalSortExpr>, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// Filter matching the state of the sort for dynamic filter pushdown. + /// If `fetch` is `Some`, this will also be set and a TopK operator may be used. + /// If `fetch` is `None`, this will be `None`. + filter: Option<Arc<RwLock<TopKDynamicFilters>>>, +} + +impl SortExec { + /// Create a new sort execution plan that produces a single, + /// sorted output partition. + pub fn new(expr: LexOrdering, input: Arc<dyn ExecutionPlan>) -> Self { + let preserve_partitioning = false; + let (cache, sort_prefix) = + Self::compute_properties(&input, expr.clone(), preserve_partitioning) + .unwrap(); + Self { + expr, + input, + metrics_set: ExecutionPlanMetricsSet::new(), + preserve_partitioning, + fetch: None, + common_sort_prefix: sort_prefix, + cache: Arc::new(cache), + filter: None, + } + } + + /// Whether this `SortExec` preserves partitioning of the children + pub fn preserve_partitioning(&self) -> bool { + self.preserve_partitioning + } + + /// Specify the partitioning behavior of this sort exec + /// + /// If `preserve_partitioning` is true, sorts each partition + /// individually, producing one sorted stream for each input partition. + /// + /// If `preserve_partitioning` is false, sorts and merges all + /// input partitions producing a single, sorted partition. + pub fn with_preserve_partitioning(mut self, preserve_partitioning: bool) -> Self { + self.preserve_partitioning = preserve_partitioning; + Arc::make_mut(&mut self.cache).partitioning = + Self::output_partitioning_helper(&self.input, self.preserve_partitioning); + if self.fetch.is_some() { + self.rebuild_filter_for_current_partitioning(); + } + self + } + + fn topk_emitter_count(&self) -> usize { + self.cache.output_partitioning().partition_count() + } + + /// Build a new shared TopK dynamic filter wrapper for this `SortExec`. + fn create_filter(&self) -> Arc<RwLock<TopKDynamicFilters>> { + let children = self + .expr + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::<Vec<_>>(); + self.create_filter_with_expr(Arc::new(DynamicFilterPhysicalExpr::new( + children, + lit(true), + ))) + } + + fn create_filter_with_expr( + &self, + expr: Arc<DynamicFilterPhysicalExpr>, + ) -> Arc<RwLock<TopKDynamicFilters>> { + Arc::new(RwLock::new( + TopKDynamicFilters::new_with_topk_emitter_count( + expr, + self.topk_emitter_count(), + ), + )) + } + + /// Rebuild the shared TopK filter wrapper for the current output partitioning. + /// + /// The dynamic filter expression is preserved, but wrapper state such as the + /// shared threshold and remaining emitter count is reset for the new + /// partitioning. + fn rebuild_filter_for_current_partitioning(&mut self) { + let filter_expr = self.filter.as_ref().map(|filter| filter.read().expr()); + if let Some(filter_expr) = filter_expr { + self.filter = Some(self.create_filter_with_expr(filter_expr)); + } + } + + fn cloned(&self) -> Self { + SortExec { + input: Arc::clone(&self.input), + expr: self.expr.clone(), + metrics_set: self.metrics_set.clone(), + preserve_partitioning: self.preserve_partitioning, + common_sort_prefix: self.common_sort_prefix.clone(), + fetch: self.fetch, + cache: Arc::clone(&self.cache), + filter: self.filter.clone(), + } + } + + /// Modify how many rows to include in the result + /// + /// If None, then all rows will be returned, in sorted order. + /// If Some, then only the top `fetch` rows will be returned. + /// This can reduce the memory pressure required by the sort + /// operation since rows that are not going to be included + /// can be dropped. + pub fn with_fetch(&self, fetch: Option<usize>) -> Self { + let mut cache = PlanProperties::clone(&self.cache); + // If the SortExec can emit incrementally (that means the sort requirements + // and properties of the input match), the SortExec can generate its result + // without scanning the entire input when a fetch value exists. + let is_pipeline_friendly = matches!( + cache.emission_type, + EmissionType::Incremental | EmissionType::Both + ); + if fetch.is_some() && is_pipeline_friendly { + cache = cache.with_boundedness(Boundedness::Bounded); + } + let mut new_sort = self.cloned(); + new_sort.fetch = fetch; + new_sort.cache = cache.into(); + if fetch.is_some() { + if new_sort.filter.is_some() { + // Keep the dynamic filter expression, but reset wrapper state + // such as the shared threshold and expected emitter count. + new_sort.rebuild_filter_for_current_partitioning(); + } else { + new_sort.filter = Some(new_sort.create_filter()); + } + } else { + new_sort.filter = None; + } + new_sort + } + + /// Input schema + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Sort expressions + pub fn expr(&self) -> &LexOrdering { + &self.expr + } + + /// If `Some(fetch)`, limits output to only the first "fetch" items + pub fn fetch(&self) -> Option<usize> { + self.fetch + } + + /// Returns the dynamic filter expression for this sort (TopK), if set. + #[deprecated( + since = "55.0.0", + note = "Use ExecutionPlan::dynamic_expressions_produced instead" + )] + pub fn dynamic_filter_expr(&self) -> Option<Arc<DynamicFilterPhysicalExpr>> { + self.filter.as_ref().map(|f| f.read().expr()) + } + + /// Replace the dynamic filter expression for this sort. + /// + /// + /// Resets any internal state which may depend on the previous dynamic filter. + /// + /// Validates that the filter's children reference valid columns in + /// the sort's input schema. + pub fn with_dynamic_filter_expr( + mut self, + filter: Arc<DynamicFilterPhysicalExpr>, + ) -> Result<Self> { + let input_schema = self.input.schema(); + for child in filter.children() { + child.data_type(&input_schema)?; + } + self.filter = Some(self.create_filter_with_expr(filter)); + Ok(self) + } + + fn output_partitioning_helper( + input: &Arc<dyn ExecutionPlan>, + preserve_partitioning: bool, + ) -> Partitioning { + // Get output partitioning: + if preserve_partitioning { + input.output_partitioning().clone() + } else { + Partitioning::UnknownPartitioning(1) + } + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + /// It also returns the common sort prefix between the input and the sort expressions. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + sort_exprs: LexOrdering, + preserve_partitioning: bool, + ) -> Result<(PlanProperties, Vec<PhysicalSortExpr>)> { + let (sort_prefix, sort_satisfied) = input + .equivalence_properties() + .extract_common_sort_prefix(sort_exprs.clone())?; + + // The emission type depends on whether the input is already sorted: + // - If already fully sorted, we can emit results in the same way as the input + // - If not sorted, we must wait until all data is processed to emit results (Final) + let emission_type = if sort_satisfied { + input.pipeline_behavior() + } else { + EmissionType::Final + }; + + // The boundedness depends on whether the input is already sorted: + // - If already sorted, we have the same property as the input + // - If not sorted and input is unbounded, we require infinite memory and generates + // unbounded data (not practical). + // - If not sorted and input is bounded, then the SortExec is bounded, too. + let boundedness = if sort_satisfied { + input.boundedness() + } else { + match input.boundedness() { + Boundedness::Unbounded { .. } => Boundedness::Unbounded { + requires_infinite_memory: true, + }, + bounded => bounded, + } + }; + + // Calculate equivalence properties; i.e. reset the ordering equivalence + // class with the new ordering: + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.reorder(sort_exprs)?; + + // Get output partitioning: + let output_partitioning = + Self::output_partitioning_helper(input, preserve_partitioning); + + Ok(( + PlanProperties::new( + eq_properties, + output_partitioning, + emission_type, + boundedness, + ), + sort_prefix, + )) + } +} + +impl DisplayAs for SortExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let preserve_partitioning = self.preserve_partitioning; + match self.fetch { + Some(fetch) => { + write!( + f, + "SortExec: TopK(fetch={fetch}), expr=[{}], preserve_partitioning=[{preserve_partitioning}]", + self.expr + )?; + if let Some(filter) = &self.filter + && let Ok(current) = filter.read().expr().current() + && !current.eq(&lit(true)) + { + write!(f, ", filter=[{current}]")?; + } + if !self.common_sort_prefix.is_empty() { + write!(f, ", sort_prefix=[")?; + let mut first = true; + for sort_expr in &self.common_sort_prefix { + if first { + first = false; + } else { + write!(f, ", ")?; + } + write!(f, "{sort_expr}")?; + } + write!(f, "]") + } else { + Ok(()) + } + } + None => write!( + f, + "SortExec: expr=[{}], preserve_partitioning=[{preserve_partitioning}]", + self.expr + ), + } + } + DisplayFormatType::TreeRender => match self.fetch { + Some(fetch) => { + writeln!(f, "{}", self.expr)?; + writeln!(f, "limit={fetch}") + } + None => { + writeln!(f, "{}", self.expr) + } + }, + } + } +} + +impl ExecutionPlan for SortExec { + fn name(&self) -> &'static str { + match self.fetch { + Some(_) => "SortExec(TopK)", + None => "SortExec", + } + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { + vec![Distribution::UnspecifiedDistribution] + } else { + // global sort + // TODO support range partitioning and OrderedDistribution. + // See https://github.com/apache/datafusion/issues/22395 + vec![Distribution::SinglePartition] + }) + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let dynamic_filter = self + .filter + .as_ref() + .map(|filter| filter.read().expr() as Arc<dyn PhysicalExpr>); + crate::apply_expression_roots( + self.expr + .iter() + .map(|sort_expr| &sort_expr.expr) + .chain(dynamic_filter.iter()), + f, + ) + } + + fn dynamic_expressions_produced(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.filter + .iter() + .map(|filter| filter.read().expr() as Arc<dyn PhysicalExpr>) + .collect() + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let mut new_sort = self.cloned(); + assert_eq!(children.len(), 1, "SortExec should have exactly one child"); + new_sort.input = Arc::clone(&children[0]); + + if options.children_properties == ChildrenPropertiesMode::Recompute { + // Recompute the properties based on the new input since they may have changed. + let (cache, sort_prefix) = Self::compute_properties( + &new_sort.input, + new_sort.expr.clone(), + new_sort.preserve_partitioning, + )?; + new_sort.cache = Arc::new(cache); + new_sort.common_sort_prefix = sort_prefix; + if new_sort.fetch.is_some() { + new_sort.rebuild_filter_for_current_partitioning(); + } + } + + Ok(Arc::new(new_sort)) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + match has_same_children_properties(self.as_ref(), &children)? { + true => self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ), + false => self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ), + } + } + + fn reset_state(self: Arc<Self>) -> Result<Arc<dyn ExecutionPlan>> { + let children = self.children().into_iter().cloned().collect(); + let new_sort = replace_children_if_necessary(self, children)?; + let mut new_sort = new_sort + .downcast_ref::<SortExec>() + .expect("rebuilt SortExec with new children") + .clone(); + // Our dynamic filter and execution metrics are the state we need to reset. + new_sort.filter = Some(new_sort.create_filter()); + new_sort.metrics_set = ExecutionPlanMetricsSet::new(); + + Ok(Arc::new(new_sort)) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start SortExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + + let mut input = self.input.execute(partition, Arc::clone(&context))?; + + let execution_options = &context.session_config().options().execution; + + trace!("End SortExec's input.execute for partition: {partition}"); + + let sort_satisfied = self + .input + .equivalence_properties() + .ordering_satisfy(self.expr.clone())?; + + match (sort_satisfied, self.fetch.as_ref()) { + (true, Some(fetch)) => Ok(Box::pin(LimitStream::new( + input, + 0, + Some(*fetch), + BaselineMetrics::new(&self.metrics_set, partition), + ))), + (true, None) => Ok(input), + (false, Some(fetch)) => { + let filter = self.filter.clone(); + let mut topk = TopK::try_new( + partition, + input.schema(), + self.common_sort_prefix.clone(), + self.expr.clone(), + *fetch, + context.session_config().batch_size(), + context.runtime_env(), + &self.metrics_set, + Arc::clone(&unwrap_or_internal_err!(filter)), + )?; + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + futures::stream::once(async move { + while let Some(batch) = input.next().await { + let batch = batch?; + topk.insert_batch(batch)?; + if topk.finished { + break; + } + } + drop(input); + topk.emit() + }) + .try_flatten(), + ))) + } + (false, None) => { + let mut sorter = ExternalSorter::new( + partition, + input.schema(), + self.expr.clone(), + context.session_config().batch_size(), + execution_options.sort_spill_reservation_bytes, + execution_options.sort_in_place_threshold_bytes, + context.session_config().spill_compression(), + &self.metrics_set, + context.runtime_env(), + )?; + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + futures::stream::once(async move { + while let Some(batch) = input.next().await { + let batch = batch?; + sorter.insert_batch(batch).await?; + } + drop(input); + sorter.sort().await + }) + .try_flatten(), + ))) + } + } + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics_set.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + let child_partition = if self.preserve_partitioning() { + partition + } else { + None + }; + vec![ChildStats::At(child_partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + Some(Arc::new(SortExec::with_fetch(self, limit))) + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + fn cardinality_effect(&self) -> CardinalityEffect { + if self.fetch.is_none() { + CardinalityEffect::Equal + } else { + CardinalityEffect::LowerEqual + } + } + + /// Tries to swap the projection with its input [`SortExec`]. If it can be done, + /// it returns the new swapped version having the [`SortExec`] as the top plan. + /// Otherwise, it returns None. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection does not narrow the schema, we should not try to push it down. + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + + let Some(updated_exprs) = update_ordering(self.expr.clone(), projection.expr())? + else { + return Ok(None); + }; + + Ok(Some(Arc::new( + SortExec::new(updated_exprs, make_with_child(projection, self.input())?) + .with_fetch(self.fetch()) + .with_preserve_partitioning(self.preserve_partitioning()), + ))) + } + + fn gather_filters_for_pushdown( + &self, + phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + config: &datafusion_common::config::ConfigOptions, + ) -> Result<FilterDescription> { + if phase != FilterPushdownPhase::Post { + if self.fetch.is_some() { + return Ok(FilterDescription::all_unsupported( + &parent_filters, + &self.children(), + )); + } + return FilterDescription::from_children(parent_filters, &self.children()); + } + + // In Post phase: block parent filters when fetch is set, + // but still push the TopK dynamic filter (self-filter). + let mut child = if self.fetch.is_some() { + ChildFilterDescription::all_unsupported(&parent_filters) + } else { + ChildFilterDescription::from_child(&parent_filters, self.input())? + }; + + if let Some(filter) = &self.filter + && config.optimizer.enable_topk_dynamic_filter_pushdown + { + child = child.with_self_filter(filter.read().expr()); + } + + Ok(FilterDescription::new().with_child(child)) + } + + fn handle_child_pushdown_result( + &self, + _phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &datafusion_common::config::ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + // For a plain sort (no fetch) we intercept any unsupported filters + // by inserting a FilterExec below this Sort. Moving the filter below + // Sort is safe because Sort preserves all rows. + // + // Why not fetch (TopK)? + // A sort with fetch limits the number of output rows. Inserting a + // FilterExec *below* the TopK would change semantics. A filter *above* + // the TopK is supposed to post-filter its output (e.g. "take the top 10 + // rows, then keep only those with a > 5"). Pushing the filter below + // Sort changes the meaning to "filter first, then take top 10", which + // produces a different result. + if self.fetch.is_some() { + return Ok(FilterPushdownPropagation::if_all(child_pushdown_result)); + } + + // Collect parent filters that were NOT successfully pushed to our child. + let unsupported_filters: Vec<Arc<dyn PhysicalExpr>> = child_pushdown_result + .parent_filters + .iter() + .filter(|&f| matches!(f.all(), PushedDown::No)) + .map(|f| Arc::clone(&f.filter)) + .collect(); + + if unsupported_filters.is_empty() { + // All filters were pushed — nothing extra to do. + return Ok(FilterPushdownPropagation::if_all(child_pushdown_result)); + } + + // Build a single conjunctive predicate from the unsupported filters + // and insert a FilterExec between this SortExec and its child. + let predicate = datafusion_physical_expr::conjunction(unsupported_filters); + let new_child = + Arc::new(FilterExec::try_new(predicate, Arc::clone(self.input()))?) + as Arc<dyn ExecutionPlan>; + let new_sort = Arc::new( + SortExec::new(self.expr.clone(), new_child) + .with_fetch(self.fetch()) + .with_preserve_partitioning(self.preserve_partitioning()), + ) as Arc<dyn ExecutionPlan>; + + Ok(FilterPushdownPropagation { + filters: vec![PushedDown::Yes; child_pushdown_result.parent_filters.len()], + updated_node: Some(new_sort), + }) + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = self + .expr() + .iter() + .map(|sort_expr| { + let sort_node = Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), + asc: !sort_expr.options.descending, + nulls_first: sort_expr.options.nulls_first, + }); + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + sort_node, + )), + }) + }) + .collect::<Result<Vec<_>>>()?; + let dynamic_filter = self + .dynamic_expressions_produced() + .into_iter() + .next() + .map(|expr| ctx.encode_expr(&expr)) + .transpose()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Sort(Box::new( + protobuf::SortExecNode { + input: Some(Box::new(input)), + expr, + fetch: match self.fetch() { + Some(n) => n as i64, + None => -1, + }, + preserve_partitioning: self.preserve_partitioning(), + dynamic_filter, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + use protobuf::physical_expr_node::ExprType; + let sort = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Sort, + "SortExec", + ); + let input = + ctx.decode_required_child(sort.input.as_deref(), "SortExec", "input")?; + let input_schema = input.schema(); + let exprs = sort + .expr + .iter() + .map(|expr| { + let Some(ExprType::Sort(sort_expr)) = expr.expr_type.as_ref() else { + return datafusion_common::internal_err!( + "SortExec expr must be a sort expression" + ); + }; + let expr_node = sort_expr.expr.as_deref().ok_or_else(|| { + internal_datafusion_err!( + "SortExec sort expression is missing its inner expr" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, input_schema.as_ref())?, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::<Result<Vec<_>>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return datafusion_common::internal_err!("SortExec requires an ordering"); + }; + let fetch = (sort.fetch >= 0).then_some(sort.fetch as usize); + let new_sort = SortExec::new(ordering, input) + .with_fetch(fetch) + .with_preserve_partitioning(sort.preserve_partitioning); + + let new_sort = if let Some(df_proto) = &sort.dynamic_filter { + let df_expr = + ctx.decode_expr(df_proto, new_sort.input().schema().as_ref())?; + let df = (df_expr as Arc<dyn std::any::Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + .map_err(|_| { + internal_datafusion_err!( + "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + new_sort.with_dynamic_filter_expr(df)? + } else { + new_sort + }; + + Ok(Arc::new(new_sort)) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::pin::Pin; + use std::task::{Context, Poll}; + + use super::*; + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::collect; + use crate::empty::EmptyExec; + use crate::execution_plan::Boundedness; + use crate::expressions::col; + use crate::filter_pushdown::{FilterPushdownPhase, PushedDown}; + use crate::test; + use crate::test::TestMemoryExec; + use crate::test::exec::{BlockingExec, assert_strong_count_converges_to_zero}; + use crate::test::{assert_is_pending, make_partition}; + + use arrow::array::*; + use arrow::compute::SortOptions; + use arrow::datatypes::*; + use datafusion_common::ScalarValue; + use datafusion_common::cast::as_primitive_array; + use datafusion_common::config::ConfigOptions; + use datafusion_common::test_util::batches_to_string; + use datafusion_execution::RecordBatchStream; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::expressions::{Column, Literal}; + use datafusion_physical_expr::{DynamicFilterTracking, EquivalenceProperties}; + + use datafusion_physical_expr_common::metrics::MetricValue; + use futures::{FutureExt, Stream, TryStreamExt}; + use insta::assert_snapshot; + + #[derive(Debug, Clone)] + pub struct SortedUnboundedExec { + schema: Schema, + batch_size: u64, + cache: Arc<PlanProperties>, + } + + impl DisplayAs for SortedUnboundedExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => write!(f, "UnboundableExec",).unwrap(), + } + Ok(()) + } + } + + impl SortedUnboundedExec { + fn compute_properties(schema: SchemaRef) -> PlanProperties { + let mut eq_properties = EquivalenceProperties::new(schema); + eq_properties.add_ordering([PhysicalSortExpr::new_default(Arc::new( + Column::new("c1", 0), + ))]); + PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ) + } + } + + impl ExecutionPlan for SortedUnboundedExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(SortedUnboundedStream { + schema: Arc::new(self.schema.clone()), + batch_size: self.batch_size, + offset: 0, + })) + } + } + + #[derive(Debug)] + pub struct SortedUnboundedStream { + schema: SchemaRef, + batch_size: u64, + offset: u64, + } + + impl Stream for SortedUnboundedStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let batch = SortedUnboundedStream::create_record_batch( + Arc::clone(&self.schema), + self.offset, + self.batch_size, + ); + self.offset += self.batch_size; + Poll::Ready(Some(Ok(batch))) + } + } + + impl RecordBatchStream for SortedUnboundedStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + } + + impl SortedUnboundedStream { + fn create_record_batch( + schema: SchemaRef, + offset: u64, + batch_size: u64, + ) -> RecordBatch { + let values = (0..batch_size).map(|i| offset + i).collect::<Vec<_>>(); + let array = UInt64Array::from(values); + let array_ref: ArrayRef = Arc::new(array); + RecordBatch::try_new(schema, vec![array_ref]).unwrap() + } + } + + #[tokio::test] + async fn test_in_mem_sort() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let partitions = 4; + let csv = test::scan_partitioned(partitions); + let schema = csv.schema(); + + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions::default(), + }] + .into(), + Arc::new(CoalescePartitionsExec::new(csv)), + )); + + let result = collect(sort_exec, Arc::clone(&task_ctx)).await?; + + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 400); + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + /// Single-column run coalescing: many small batches above a tiny in-place + /// threshold (with ample memory, so no spill) must still produce a correct + /// total order, including NULLs. + #[tokio::test] + async fn test_in_mem_sort_coalesced_runs() -> Result<()> { + // Tiny in-place threshold forces the sort-then-merge path and, for a + // single column, the coalescing branch. Ample memory => no spill. + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(64) + .with_sort_in_place_threshold_bytes(1024), + ), + ); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + // Build many small batches of shuffled values with interspersed NULLs, + // so coalescing produces several multi-row runs that must be merged. + let num_batches = 40; + let rows_per_batch = 50; + let mut all_values: Vec<Option<i32>> = Vec::new(); + let mut batches = Vec::with_capacity(num_batches); + for b in 0..num_batches { + let mut col_values: Vec<Option<i32>> = Vec::with_capacity(rows_per_batch); + for r in 0..rows_per_batch { + let idx = (b * rows_per_batch + r) as i64; + // Deterministic scramble to avoid any pre-existing ordering. + let scrambled = ((idx.wrapping_mul(2_654_435_761)) % 1000) as i32; + let v = if idx % 7 == 0 { None } else { Some(scrambled) }; + col_values.push(v); + all_values.push(v); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(col_values))], + )?; + batches.push(batch); + } + let total_rows = num_batches * rows_per_batch; + + let options = SortOptions::default(); + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("a", &schema)?, + options, + }] + .into(), + TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?, + )); + + let result = collect( + Arc::clone(&sort_exec) as Arc<dyn ExecutionPlan>, + Arc::clone(&task_ctx), + ) + .await?; + + // Flatten the sorted output. + let mut got: Vec<Option<i32>> = Vec::with_capacity(total_rows); + for batch in &result { + let arr = as_primitive_array::<Int32Type>(batch.column(0))?; + for i in 0..arr.len() { + got.push(if arr.is_null(i) { + None + } else { + Some(arr.value(i)) + }); + } + } + assert_eq!(got.len(), total_rows, "row count must be preserved"); + + // Reference: sort the original values with the same semantics + // (ascending, NULLs first per SortOptions::default()). + let mut expected = all_values.clone(); + expected.sort_by(|a, b| match (a, b) { + (None, None) => std::cmp::Ordering::Equal, + (None, Some(_)) => std::cmp::Ordering::Less, // nulls_first + (Some(_), None) => std::cmp::Ordering::Greater, + (Some(x), Some(y)) => x.cmp(y), + }); + + assert_eq!( + got, expected, + "coalesced-run sort output must be totally ordered" + ); + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_sort_spill() -> Result<()> { + // trigger spill w/ 100 batches + let session_config = SessionConfig::new(); + let sort_spill_reservation_bytes = session_config + .options() + .execution + .sort_spill_reservation_bytes; + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(sort_spill_reservation_bytes + 12288, 1.0) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + // The input has 100 partitions, each partition has a batch containing 100 rows. + // Each row has a single Int32 column with values 0..100. The total size of the + // input is roughly 40000 bytes. + let partitions = 100; + let input = test::scan_partitioned(partitions); + let schema = input.schema(); + + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions::default(), + }] + .into(), + Arc::new(CoalescePartitionsExec::new(input)), + )); + + let result = collect( + Arc::clone(&sort_exec) as Arc<dyn ExecutionPlan>, + Arc::clone(&task_ctx), + ) + .await?; + + assert_eq!(result.len(), 2); + + // Now, validate metrics + let metrics = sort_exec.metrics().unwrap(); + + assert_eq!(metrics.output_rows().unwrap(), 10000); + assert!(metrics.elapsed_compute().unwrap() > 0); + + let spill_count = metrics.spill_count().unwrap(); + let spilled_rows = metrics.spilled_rows().unwrap(); + let spilled_bytes = metrics.spilled_bytes().unwrap(); + // Processing 40000 bytes of data using 12288 bytes of memory requires 3 spills + // unless we do something really clever. It will spill roughly 9000+ rows and 36000 + // bytes. We leave a little wiggle room for the actual numbers. + assert!((3..=10).contains(&spill_count)); + assert!((9000..=10000).contains(&spilled_rows)); + assert!((38000..=44000).contains(&spilled_bytes)); + + let columns = result[0].columns(); + + let i = as_primitive_array::<Int32Type>(&columns[0])?; + assert_eq!(i.value(0), 0); + assert_eq!(i.value(i.len() - 1), 81); + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_batch_reservation_error() -> Result<()> { + // Pick a memory limit and sort_spill_reservation that make the first batch reservation fail. + let merge_reservation: usize = 0; // Set to 0 for simplicity + + let session_config = + SessionConfig::new().with_sort_spill_reservation_bytes(merge_reservation); + + let plan = test::scan_partitioned(1); + + // Read the first record batch to determine the actual memory requirement + let expected_batch_reservation = { + let temp_ctx = Arc::new(TaskContext::default()); + let mut stream = plan.execute(0, Arc::clone(&temp_ctx))?; + let first_batch = stream.next().await.unwrap()?; + get_reserved_bytes_for_record_batch(&first_batch)? + }; + + // Set memory limit just short of what we need + let memory_limit: usize = expected_batch_reservation + merge_reservation - 1; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(memory_limit, 1.0) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + // Verify that our memory limit is insufficient + { + let mut stream = plan.execute(0, Arc::clone(&task_ctx))?; + let first_batch = stream.next().await.unwrap()?; + let batch_reservation = get_reserved_bytes_for_record_batch(&first_batch)?; + + assert_eq!(batch_reservation, expected_batch_reservation); + assert!(memory_limit < (merge_reservation + batch_reservation)); + } + + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr::new_default(col("i", &plan.schema())?)].into(), + plan, + )); + + let result = collect(Arc::clone(&sort_exec) as _, Arc::clone(&task_ctx)).await; + + let err = result.unwrap_err(); + assert!( + matches!(err, DataFusionError::Context(..)), + "Assertion failed: expected a Context error, but got: {err:?}" + ); + + // Assert that the context error is wrapping a resources exhausted error. + assert!( + matches!(err.find_root(), DataFusionError::ResourcesExhausted(_)), + "Assertion failed: expected a ResourcesExhausted error, but got: {err:?}" + ); + + // Verify external sorter error message when resource is exhausted + let config_vector = vec![ + "datafusion.runtime.memory_limit", + "datafusion.execution.sort_spill_reservation_bytes", + ]; + let error_message = err.message().to_string(); + for config in config_vector.into_iter() { + assert!( + error_message.as_str().contains(config), + "Config: '{}' should be contained in error message: {}.", + config, + error_message.as_str() + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_sort_spill_utf8_strings() -> Result<()> { + let session_config = SessionConfig::new() + .with_batch_size(100) + .with_sort_in_place_threshold_bytes(20 * 1024) + .with_sort_spill_reservation_bytes(100 * 1024); + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(500 * 1024, 1.0) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), + ); + + // The input has 200 partitions, each partition has a batch containing 100 rows. + // Each row has a single Utf8 column, the Utf8 string values are roughly 42 bytes. + // The total size of the input is roughly 820 KB. + let input = test::scan_partitioned_utf8(200); + let schema = input.schema(); + + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions::default(), + }] + .into(), + Arc::new(CoalescePartitionsExec::new(input)), + )); + + let result = collect(Arc::clone(&sort_exec) as _, Arc::clone(&task_ctx)).await?; + + let num_rows = result.iter().map(|batch| batch.num_rows()).sum::<usize>(); + assert_eq!(num_rows, 20000); + + // Now, validate metrics + let metrics = sort_exec.metrics().unwrap(); + + assert_eq!(metrics.output_rows().unwrap(), 20000); + assert!(metrics.elapsed_compute().unwrap() > 0); + + let spill_count = metrics.spill_count().unwrap(); + let spilled_rows = metrics.spilled_rows().unwrap(); + let spilled_bytes = metrics.spilled_bytes().unwrap(); + + // This test case is processing 840KB of data using 400KB of memory. Note + // that buffered batches can't be dropped until all sorted batches are + // generated, so we can only buffer `sort_spill_reservation_bytes` of sorted + // batches. + // The number of spills is roughly calculated as: + // `number_of_batches / (sort_spill_reservation_bytes / batch_size)` + + // If this assertion fail with large spill count, make sure the following + // case does not happen: + // During external sorting, one sorted run should be spilled to disk in a + // single file, due to memory limit we might need to append to the file + // multiple times to spill all the data. Make sure we're not writing each + // appending as a separate file. + assert!((4..=8).contains(&spill_count)); + assert!((15000..=20000).contains(&spilled_rows)); + assert!((900000..=1000000).contains(&spilled_bytes)); + + // Verify that the result is sorted + let concated_result = concat_batches(&schema, &result)?; + let columns = concated_result.columns(); + let string_array = as_string_array(&columns[0]); + for i in 0..string_array.len() - 1 { + assert!(string_array.value(i) <= string_array.value(i + 1)); + } + + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_sort_fetch_memory_calculation() -> Result<()> { + // This test mirrors down the size from the example above. + let avg_batch_size = 400; + let partitions = 4; + + // A tuple of (fetch, expect_spillage) + let test_options = vec![ + // Since we don't have a limit (and the memory is less than the total size of + // all the batches we are processing, we expect it to spill. + (None, true), + // When we have a limit however, the buffered size of batches should fit in memory + // since it is much lower than the total size of the input batch. + (Some(1), false), + ]; + + for (fetch, expect_spillage) in test_options { + let session_config = SessionConfig::new(); + let sort_spill_reservation_bytes = session_config + .options() + .execution + .sort_spill_reservation_bytes; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit( + sort_spill_reservation_bytes + avg_batch_size * (partitions - 1), + 1.0, + ) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(session_config), + ); + + let csv = test::scan_partitioned(partitions); + let schema = csv.schema(); + + let sort_exec = Arc::new( + SortExec::new( + [PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions::default(), + }] + .into(), + Arc::new(CoalescePartitionsExec::new(csv)), + ) + .with_fetch(fetch), + ); + + let result = + collect(Arc::clone(&sort_exec) as _, Arc::clone(&task_ctx)).await?; + assert_eq!(result.len(), 1); + + let metrics = sort_exec.metrics().unwrap(); + let did_it_spill = metrics.spill_count().unwrap_or(0) > 0; + assert_eq!(did_it_spill, expect_spillage, "with fetch: {fetch:?}"); + } + Ok(()) + } + + #[tokio::test] + async fn test_sort_memory_reduction_per_batch() -> Result<()> { + // This test verifies that memory reservation is reduced for every batch emitted + // during the sort process. This is important to ensure we don't hold onto + // memory longer than necessary. + + // Create a large enough batch that will be split into multiple output batches + let batch_size = 50; // Small batch size to force multiple output batches + let num_rows = 1000; // Create enough data for multiple batches + + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(batch_size) + .with_sort_in_place_threshold_bytes(usize::MAX), // Ensure we don't concat batches + ), + ); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create unsorted data + let mut values: Vec<i32> = (0..num_rows).collect(); + values.reverse(); + + let input_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + )?; + + let batches = vec![input_batch]; + + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: Arc::new(Column::new("a", 0)), + options: SortOptions::default(), + }] + .into(), + TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?, + )); + + let mut stream = sort_exec.execute(0, Arc::clone(&task_ctx))?; + + let mut previous_reserved = task_ctx.runtime_env().memory_pool.reserved(); + let mut batch_count = 0; + + // Collect batches and verify memory is reduced with each batch + while let Some(result) = stream.next().await { + let batch = result?; + batch_count += 1; + + // Verify we got a non-empty batch + assert!(batch.num_rows() > 0, "Batch should not be empty"); + + let current_reserved = task_ctx.runtime_env().memory_pool.reserved(); + + // After the first batch, memory should be reducing or staying the same + // (it should not increase as we emit batches) + if batch_count > 1 { + assert!( + current_reserved <= previous_reserved, + "Memory reservation should decrease or stay same as batches are emitted. \ + Batch {batch_count}: previous={previous_reserved}, current={current_reserved}" + ); + } + + previous_reserved = current_reserved; + } + + assert!( + batch_count > 1, + "Expected multiple batches to be emitted, got {batch_count}" + ); + + // Verify all memory is returned at the end + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "All memory should be returned after consuming all batches" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_sort_metadata() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let field_metadata: HashMap<String, String> = + vec![("foo".to_string(), "bar".to_string())] + .into_iter() + .collect(); + let schema_metadata: HashMap<String, String> = + vec![("baz".to_string(), "barf".to_string())] + .into_iter() + .collect(); + + let mut field = Field::new("field_name", DataType::UInt64, true); + field.set_metadata(field_metadata.clone()); + let schema = Schema::new_with_metadata(vec![field], schema_metadata.clone()); + let schema = Arc::new(schema); + + let data: ArrayRef = + Arc::new(vec![3, 2, 1].into_iter().map(Some).collect::<UInt64Array>()); + + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![data])?; + let input = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("field_name", &schema)?, + options: SortOptions::default(), + }] + .into(), + input, + )); + + let result: Vec<RecordBatch> = collect(sort_exec, task_ctx).await?; + + let expected_data: ArrayRef = + Arc::new(vec![1, 2, 3].into_iter().map(Some).collect::<UInt64Array>()); + let expected_batch = + RecordBatch::try_new(Arc::clone(&schema), vec![expected_data])?; + + // Data is correct + assert_eq!(&vec![expected_batch], &result); + + // explicitly ensure the metadata is present + assert_eq!(result[0].schema().fields()[0].metadata(), &field_metadata); + assert_eq!(result[0].schema().metadata(), &schema_metadata); + + Ok(()) + } + + #[tokio::test] + async fn test_lex_sort_by_mixed_types() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "b", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + ])); + + // define data. + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(2), None, Some(1), Some(2)])), + Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ + Some(vec![Some(3)]), + Some(vec![Some(1)]), + Some(vec![Some(6), None]), + Some(vec![Some(5)]), + ])), + ], + )?; + + let sort_exec = Arc::new(SortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + ] + .into(), + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?, + )); + + assert_eq!(DataType::Int32, *sort_exec.schema().field(0).data_type()); + assert_eq!( + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + *sort_exec.schema().field(1).data_type() + ); + + let result: Vec<RecordBatch> = + collect(Arc::clone(&sort_exec) as Arc<dyn ExecutionPlan>, task_ctx).await?; + let metrics = sort_exec.metrics().unwrap(); + assert!(metrics.elapsed_compute().unwrap() > 0); + assert_eq!(metrics.output_rows().unwrap(), 4); + assert_eq!(result.len(), 1); + + let expected = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![None, Some(1), Some(2), Some(2)])), + Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ + Some(vec![Some(1)]), + Some(vec![Some(6), None]), + Some(vec![Some(5)]), + Some(vec![Some(3)]), + ])), + ], + )?; + + assert_eq!(expected, result[0]); + + Ok(()) + } + + #[tokio::test] + async fn test_lex_sort_by_float() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float32, true), + Field::new("b", DataType::Float64, true), + ])); + + // define data. + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Float32Array::from(vec![ + Some(f32::NAN), + None, + None, + Some(f32::NAN), + Some(1.0_f32), + Some(1.0_f32), + Some(2.0_f32), + Some(3.0_f32), + ])), + Arc::new(Float64Array::from(vec![ + Some(200.0_f64), + Some(20.0_f64), + Some(10.0_f64), + Some(100.0_f64), + Some(f64::NAN), + None, + None, + Some(f64::NAN), + ])), + ], + )?; + + let sort_exec = Arc::new(SortExec::new( + [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: true, + nulls_first: true, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }, + ] + .into(), + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)?, + )); + + assert_eq!(DataType::Float32, *sort_exec.schema().field(0).data_type()); + assert_eq!(DataType::Float64, *sort_exec.schema().field(1).data_type()); + + let result: Vec<RecordBatch> = + collect(Arc::clone(&sort_exec) as Arc<dyn ExecutionPlan>, task_ctx).await?; + let metrics = sort_exec.metrics().unwrap(); + assert!(metrics.elapsed_compute().unwrap() > 0); + assert_eq!(metrics.output_rows().unwrap(), 8); + assert_eq!(result.len(), 1); + + let columns = result[0].columns(); + + assert_eq!(DataType::Float32, *columns[0].data_type()); + assert_eq!(DataType::Float64, *columns[1].data_type()); + + let a = as_primitive_array::<Float32Type>(&columns[0])?; + let b = as_primitive_array::<Float64Type>(&columns[1])?; + + // convert result to strings to allow comparing to expected result containing NaN + let result: Vec<(Option<String>, Option<String>)> = (0..result[0].num_rows()) + .map(|i| { + let aval = if a.is_valid(i) { + Some(a.value(i).to_string()) + } else { + None + }; + let bval = if b.is_valid(i) { + Some(b.value(i).to_string()) + } else { + None + }; + (aval, bval) + }) + .collect(); + + let expected: Vec<(Option<String>, Option<String>)> = vec![ + (None, Some("10".to_owned())), + (None, Some("20".to_owned())), + (Some("NaN".to_owned()), Some("100".to_owned())), + (Some("NaN".to_owned()), Some("200".to_owned())), + (Some("3".to_owned()), Some("NaN".to_owned())), + (Some("2".to_owned()), None), + (Some("1".to_owned()), Some("NaN".to_owned())), + (Some("1".to_owned()), None), + ]; + + assert_eq!(expected, result); + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions::default(), + }] + .into(), + blocking_exec, + )); + + let fut = collect(sort_exec, Arc::clone(&task_ctx)); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + + #[test] + fn test_empty_sort_batch() { + let schema = Arc::new(Schema::empty()); + let options = RecordBatchOptions::new().with_row_count(Some(1)); + let batch = + RecordBatch::try_new_with_options(Arc::clone(&schema), vec![], &options) + .unwrap(); + + let expressions = [PhysicalSortExpr { + expr: Arc::new(Literal::new(ScalarValue::Int64(Some(1)))), + options: SortOptions::default(), + }] + .into(); + + let result = sort_batch(&batch, &expressions, None).unwrap(); + assert_eq!(result.num_rows(), 1); + } + + #[tokio::test] + async fn topk_unbounded_source() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); + let source = SortedUnboundedExec { + schema: schema.clone(), + batch_size: 2, + cache: Arc::new(SortedUnboundedExec::compute_properties(Arc::new( + schema.clone(), + ))), + }; + let mut plan = SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new( + "c1", 0, + )))] + .into(), + Arc::new(source), + ); + plan = plan.with_fetch(Some(9)); + + let batches = collect(Arc::new(plan), task_ctx).await?; + assert_snapshot!(batches_to_string(&batches), @r" + +----+ + | c1 | + +----+ + | 0 | + | 1 | + | 2 | + | 3 | + | 4 | + | 5 | + | 6 | + | 7 | + | 8 | + +----+ + "); + Ok(()) + } + + #[tokio::test] + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics() + -> Result<()> { + let batch_size = 100; + + let create_task_ctx = |_: &[RecordBatch]| { + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(batch_size) + .with_sort_in_place_threshold_bytes(usize::MAX), + ) + }; + + // Smaller than batch size and require more than a single batch to get the requested batch size + test_sort_output_batch_size_and_base_metrics(10, batch_size / 4, create_task_ctx) + .await?; + + // Not evenly divisible by batch size + test_sort_output_batch_size_and_base_metrics(10, batch_size + 7, create_task_ctx) + .await?; + + // Evenly divisible by batch size and is larger than 2 output batches + test_sort_output_batch_size_and_base_metrics(10, batch_size * 3, create_task_ctx) + .await?; + + Ok(()) + } + + #[tokio::test] + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_sorting_in_place() + -> Result<()> { + let batch_size = 100; + + let create_task_ctx = |_: &[RecordBatch]| { + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(batch_size) + .with_sort_in_place_threshold_bytes(usize::MAX - 1), + ) + }; + + // Smaller than batch size and require more than a single batch to get the requested batch size + { + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size / 4, + create_task_ctx, + ) + .await?; + + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spills when sorting in place" + ); + } + + // Not evenly divisible by batch size + { + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size + 7, + create_task_ctx, + ) + .await?; + + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spills when sorting in place" + ); + } + + // Evenly divisible by batch size and is larger than 2 output batches + { + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size * 3, + create_task_ctx, + ) + .await?; + + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spills when sorting in place" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_a_single_batch() + -> Result<()> { + let batch_size = 100; + + let create_task_ctx = |_: &[RecordBatch]| { + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(batch_size)) + }; + + // Smaller than batch size and require more than a single batch to get the requested batch size + { + let metrics = test_sort_output_batch_size_and_base_metrics( + // Single batch + 1, + batch_size / 4, + create_task_ctx, + ) + .await?; + + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spills when sorting in place" + ); + } + + // Not evenly divisible by batch size + { + let metrics = test_sort_output_batch_size_and_base_metrics( + // Single batch + 1, + batch_size + 7, + create_task_ctx, + ) + .await?; + + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spills when sorting in place" + ); + } + + // Evenly divisible by batch size and is larger than 2 output batches + { + let metrics = test_sort_output_batch_size_and_base_metrics( + // Single batch + 1, + batch_size * 3, + create_task_ctx, + ) + .await?; + + assert_eq!( + metrics.spill_count(), + Some(0), + "Expected no spills when sorting in place" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_to_spill() + -> Result<()> { + let batch_size = 100; + + let create_task_ctx = |generated_batches: &[RecordBatch]| { + let batches_memory = generated_batches + .iter() + .map(|b| b.get_array_memory_size()) + .sum::<usize>(); + + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(batch_size) + // To make sure there is no in place sorting + .with_sort_in_place_threshold_bytes(1) + .with_sort_spill_reservation_bytes(1), + ) + .with_runtime( + RuntimeEnvBuilder::default() + .with_memory_limit(batches_memory, 1.0) + .build_arc() + .unwrap(), + ) + }; + + // Smaller than batch size and require more than a single batch to get the requested batch size + { + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size / 4, + create_task_ctx, + ) + .await?; + + assert_ne!(metrics.spill_count().unwrap(), 0, "expected to spill"); + } + + // Not evenly divisible by batch size + { + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size + 7, + create_task_ctx, + ) + .await?; + + assert_ne!(metrics.spill_count().unwrap(), 0, "expected to spill"); + } + + // Evenly divisible by batch size and is larger than 2 batches + { + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size * 3, + create_task_ctx, + ) + .await?; + + assert_ne!(metrics.spill_count().unwrap(), 0, "expected to spill"); + } + + Ok(()) + } + + async fn test_sort_output_batch_size_and_base_metrics( + number_of_batches: usize, + batch_size_to_generate: usize, + create_task_ctx: impl Fn(&[RecordBatch]) -> TaskContext, + ) -> Result<MetricsSet> { + let batches = (0..number_of_batches) + .map(|_| make_partition(batch_size_to_generate as i32)) + .collect::<Vec<_>>(); + let task_ctx = create_task_ctx(batches.as_slice()); + + let output_rows = batches.iter().map(|item| item.num_rows()).sum(); + + let expected_batch_size = task_ctx.session_config().batch_size(); + + let schema = batches[0].schema(); + let (mut output_batches, metrics) = + run_sort_on_input(task_ctx, "i", batches, schema).await?; + + let last_batch = output_batches.pop().unwrap(); + + for batch in output_batches { + assert_eq!(batch.num_rows(), expected_batch_size); + } + + let mut last_expected_batch_size = + (batch_size_to_generate * number_of_batches) % expected_batch_size; + if last_expected_batch_size == 0 { + last_expected_batch_size = expected_batch_size; + } + assert_eq!(last_batch.num_rows(), last_expected_batch_size); + + assert_baseline_metrics_for_non_empty_output( + &metrics, + output_rows, + expected_batch_size, + ); + + Ok(metrics) + } + + #[tokio::test] + async fn empty_sort_stream_should_report_end_time() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let task_ctx = TaskContext::default(); + + let (_, metrics) = run_sort_on_input(task_ctx, "i", vec![], schema).await?; + + let end_time = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::EndTimestamp(end) => Some(end), + _ => None, + }) + .expect("Must have end time metric since it exists in the baseline"); + + assert_eq!( + metrics.spill_count().unwrap_or_default(), + 0, + "expected to not have spills" + ); + assert_ne!(end_time.value(), None); + + Ok(()) + } + + fn assert_baseline_metrics_for_non_empty_output( + metrics: &MetricsSet, + output_rows: usize, + batch_size: usize, + ) { + let end_time = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::EndTimestamp(end) => Some(end), + _ => None, + }) + .expect("Must have end time metric since it exists in the baseline"); + + assert_ne!(end_time.value(), None); + + assert_eq!(metrics.output_rows(), Some(output_rows)); + + let output_bytes = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::OutputBytes(total) => Some(total), + _ => None, + }) + .expect("Must have output_bytes metric since it exists in the baseline"); + + assert_ne!(output_bytes.value(), 0_usize); + + let output_batches = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::OutputBatches(total) => Some(total), + _ => None, + }) + .expect("Must have output_batches metric since it exists in the baseline"); + + assert_eq!(output_batches.value(), output_rows.div_ceil(batch_size)); + } + + async fn run_sort_on_input( + task_ctx: TaskContext, + order_by_col: &str, + batches: Vec<RecordBatch>, + schema: SchemaRef, + ) -> Result<(Vec<RecordBatch>, MetricsSet)> { + let task_ctx = Arc::new(task_ctx); + + // let task_ctx = env. + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col(order_by_col, &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let sort_exec: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new( + ordering.clone(), + TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?, + )); + + let sorted_batches = + collect(Arc::clone(&sort_exec), Arc::clone(&task_ctx)).await?; + + let metrics = sort_exec.metrics().expect("sort have metrics"); + + // assert output + { + let input_batches_concat = concat_batches(&schema, &batches)?; + let sorted_input_batch = sort_batch(&input_batches_concat, &ordering, None)?; + + let sorted_batches_concat = concat_batches(&schema, &sorted_batches)?; + + assert_eq!(sorted_input_batch, sorted_batches_concat); + } + + Ok((sorted_batches, metrics)) + } + + #[tokio::test] + async fn test_sort_batch_chunked_basic() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create a batch with 1000 rows + let mut values: Vec<i32> = (0..1000).collect(); + // Shuffle to make it unsorted + values.reverse(); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + )?; + + let expressions: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + // Sort with batch_size = 250 + let result_batches = sort_batch_chunked(&batch, &expressions, 250)?; + + // Verify 4 batches are returned + assert_eq!(result_batches.len(), 4); + + // Verify each batch has <= 250 rows + let mut total_rows = 0; + for (i, batch) in result_batches.iter().enumerate() { + assert!( + batch.num_rows() <= 250, + "Batch {} has {} rows, expected <= 250", + i, + batch.num_rows() + ); + total_rows += batch.num_rows(); + } + + // Verify total row count matches input + assert_eq!(total_rows, 1000); + + // Verify data is correctly sorted across all chunks + let concatenated = concat_batches(&schema, &result_batches)?; + let array = as_primitive_array::<Int32Type>(concatenated.column(0))?; + for i in 0..array.len() - 1 { + assert!( + array.value(i) <= array.value(i + 1), + "Array not sorted at position {}: {} > {}", + i, + array.value(i), + array.value(i + 1) + ); + } + assert_eq!(array.value(0), 0); + assert_eq!(array.value(array.len() - 1), 999); + + Ok(()) + } + + #[tokio::test] + async fn test_sort_batch_chunked_smaller_than_batch_size() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create a batch with 50 rows + let values: Vec<i32> = (0..50).rev().collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + )?; + + let expressions: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + // Sort with batch_size = 100 + let result_batches = sort_batch_chunked(&batch, &expressions, 100)?; + + // Should return exactly 1 batch + assert_eq!(result_batches.len(), 1); + assert_eq!(result_batches[0].num_rows(), 50); + + // Verify it's correctly sorted + let array = as_primitive_array::<Int32Type>(result_batches[0].column(0))?; + for i in 0..array.len() - 1 { + assert!(array.value(i) <= array.value(i + 1)); + } + assert_eq!(array.value(0), 0); + assert_eq!(array.value(49), 49); + + Ok(()) + } + + #[tokio::test] + async fn test_sort_batch_chunked_exact_multiple() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create a batch with 1000 rows + let values: Vec<i32> = (0..1000).rev().collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + )?; + + let expressions: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + // Sort with batch_size = 100 + let result_batches = sort_batch_chunked(&batch, &expressions, 100)?; + + // Should return exactly 10 batches of 100 rows each + assert_eq!(result_batches.len(), 10); + for batch in &result_batches { + assert_eq!(batch.num_rows(), 100); + } + + // Verify sorted correctly across all batches + let concatenated = concat_batches(&schema, &result_batches)?; + let array = as_primitive_array::<Int32Type>(concatenated.column(0))?; + for i in 0..array.len() - 1 { + assert!(array.value(i) <= array.value(i + 1)); + } + + Ok(()) + } + + #[tokio::test] + async fn test_sort_batch_chunked_empty_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + let batch = RecordBatch::new_empty(Arc::clone(&schema)); + + let expressions: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + let result_batches = sort_batch_chunked(&batch, &expressions, 100)?; + + // Empty input produces no output batches (0 chunks) + assert_eq!(result_batches.len(), 0); + + Ok(()) + } + + #[tokio::test] + async fn test_get_reserved_bytes_for_record_batch_with_sliced_batches() -> Result<()> + { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create a larger batch then slice it + let large_array = Int32Array::from((0..1000).collect::<Vec<i32>>()); + let sliced_array = large_array.slice(100, 50); // Take 50 elements starting at 100 + + let sliced_batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(sliced_array)])?; + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(large_array)])?; + + let sliced_reserved = get_reserved_bytes_for_record_batch(&sliced_batch)?; + let reserved = get_reserved_bytes_for_record_batch(&batch)?; + + // The reserved memory for the sliced batch should be less than that of the full batch + assert!(reserved > sliced_reserved); + + Ok(()) + } + + #[test] + fn test_with_dynamic_filter() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let sort = SortExec::new( + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::new(Column::new("a", 0)), + options: SortOptions::default(), + }]) + .unwrap(), + child, + ) + .with_fetch(Some(10)); + + // SortExec with fetch creates a dynamic filter automatically. + let produced = sort.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + let original_id = produced[0] + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"); + + // with_dynamic_filter replaces it with a new TopKDynamicFilters. + let new_df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0)) as _], + lit(true), + )); + let new_id = new_df + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"); + let sort = sort.with_dynamic_filter_expr(Arc::clone(&new_df))?; + let produced = sort.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + let restored_id = produced[0] + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"); + assert_eq!(restored_id, new_id); + assert_ne!(restored_id, original_id); + Ok(()) + } + + async fn emit_sort_partition( + sort: &Arc<SortExec>, + partition: usize, + task_ctx: Arc<TaskContext>, + ) -> Result<()> { + let _batches: Vec<RecordBatch> = + sort.execute(partition, task_ctx)?.try_collect().await?; + Ok(()) + } + + fn assert_filter_still_waiting(filter: &Arc<DynamicFilterPhysicalExpr>) { + let dynamic_filter_expr: Arc<dyn PhysicalExpr> = + Arc::<DynamicFilterPhysicalExpr>::clone(filter); + assert!( + matches!( + DynamicFilterTracking::classify(&dynamic_filter_expr), + DynamicFilterTracking::Watching(_) + ), + "the shared filter should remain watchable until every partition emits" + ); + } + + fn dynamic_filter_produced( + plan: &dyn ExecutionPlan, + ) -> Arc<DynamicFilterPhysicalExpr> { + let expr = plan + .dynamic_expressions_produced() + .into_iter() + .next() + .expect("plan should produce a dynamic filter"); + (expr as Arc<dyn std::any::Any + Send + Sync>) + .downcast::<DynamicFilterPhysicalExpr>() + .expect("produced expression should be a DynamicFilterPhysicalExpr") + } + + #[tokio::test] + async fn test_preserved_topk_filter_waits_for_all_sort_partitions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = vec![ + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![3, 1, 2]))], + )?], + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![6, 4, 5]))], + )?], + ]; + let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let sort = SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(), + input, + ) + // `with_fetch` creates the TopK filter; preserving partitioning after + // that must rebuild it with one emitter per output partition. + .with_fetch(Some(2)) + .with_preserve_partitioning(true); + + let dynamic_filter = dynamic_filter_produced(&sort); + let sort = Arc::new(sort); + let task_ctx = Arc::new(TaskContext::default()); + + emit_sort_partition(&sort, 0, Arc::clone(&task_ctx)).await?; + assert_filter_still_waiting(&dynamic_filter); + + emit_sort_partition(&sort, 1, task_ctx).await?; + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter.wait_complete(), + ) + .await + .expect("the final preserved SortExec partition should complete the filter"); + + Ok(()) + } + + #[tokio::test] + async fn test_with_fetch_rebuilds_existing_topk_filter() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = vec![ + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![3, 1, 2]))], + )?], + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![6, 4, 5]))], + )?], + ]; + let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0))], + lit(true), + )); + let dynamic_filter_id = dynamic_filter + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"); + let sort = SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(), + input, + ) + .with_dynamic_filter_expr(dynamic_filter)? + .with_preserve_partitioning(true) + .with_fetch(Some(2)); + + let dynamic_filter = dynamic_filter_produced(&sort); + assert_eq!( + dynamic_filter + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"), + dynamic_filter_id + ); + + let sort = Arc::new(sort); + let task_ctx = Arc::new(TaskContext::default()); + + emit_sort_partition(&sort, 0, Arc::clone(&task_ctx)).await?; + assert_filter_still_waiting(&dynamic_filter); + + emit_sort_partition(&sort, 1, task_ctx).await?; + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter.wait_complete(), + ) + .await + .expect("the final preserved SortExec partition should complete the filter"); + + Ok(()) + } + + #[test] + fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + + let sort = SortExec::new( + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::new(Column::new("a", 0)), + options: SortOptions::default(), + }]) + .unwrap(), + child, + ) + .with_fetch(Some(10)); + + // Column index 99 is out of bounds for the input schema. + let df = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("bad", 99)) as _], + lit(true), + )); + assert!(sort.with_dynamic_filter_expr(df).is_err()); + Ok(()) + } + + /// Verifies that `ExternalSorter::sort()` transfers the pre-reserved + /// merge bytes to the merge stream via `take()`, rather than leaving + /// them in the sorter (via `new_empty()`). + /// + /// 1. Create a sorter with a tight memory pool and insert enough data + /// to force spilling + /// 2. Verify `merge_reservation` holds the pre-reserved bytes before sort + /// 3. Call `sort()` to get the merge stream + /// 4. Verify `merge_reservation` is now 0 (bytes transferred to merge stream) + /// 5. Simulate contention: a competing consumer grabs all available pool memory + /// 6. Verify the merge stream still works (it uses its pre-reserved bytes + /// as initial budget, not requesting from pool starting at 0) + /// + /// With `new_empty()` (before fix), step 4 fails: `merge_reservation` + /// still holds the bytes, the merge stream starts with 0 budget, and + /// those bytes become unaccounted-for reserved memory that nobody uses. + #[tokio::test] + async fn test_sort_merge_reservation_transferred_not_freed() -> Result<()> { + let sort_spill_reservation_bytes: usize = 10 * 1024; // 10 KB + + // Pool: merge reservation (10KB) + enough room for sort to work. + // The room must accommodate batch data accumulation before spilling. + let sort_working_memory: usize = 40 * 1024; // 40 KB for sort operations + let pool_size = sort_spill_reservation_bytes + sort_working_memory; + let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(pool_size)); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + + let metrics_set = ExecutionPlanMetricsSet::new(); + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); + + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(), + 128, // batch_size + sort_spill_reservation_bytes, + usize::MAX, // sort_in_place_threshold_bytes (high to avoid concat path) + SpillCompression::Uncompressed, + &metrics_set, + Arc::clone(&runtime), + )?; + + // Insert enough data to force spilling. + let num_batches = 200; + for i in 0..num_batches { + let values: Vec<i32> = ((i * 100)..((i + 1) * 100)).rev().collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(values))], + )?; + sorter.insert_batch(batch).await?; + } + + assert!( + sorter.spilled_before(), + "Test requires spilling to exercise the merge path" + ); + + // Before sort(), merge_reservation holds sort_spill_reservation_bytes. + assert!( + sorter.merge_reservation_size() >= sort_spill_reservation_bytes, + "merge_reservation should hold the pre-reserved bytes before sort()" + ); + + // Call sort() to get the merge stream. With the fix (take()), + // the pre-reserved merge bytes are transferred to the merge + // stream. Without the fix (free() + new_empty()), the bytes + // are released back to the pool and the merge stream starts + // with 0 bytes. + let merge_stream = sorter.sort().await?; + + // THE KEY ASSERTION: after sort(), merge_reservation must be 0. + // This proves take() transferred the bytes to the merge stream, + // rather than them being freed back to the pool where other + // partitions could steal them. + assert_eq!( + sorter.merge_reservation_size(), + 0, + "After sort(), merge_reservation should be 0 (bytes transferred \ + to merge stream via take()). If non-zero, the bytes are still \ + held by the sorter and will be freed on drop, allowing other \ + partitions to steal them." + ); + + // Drop the sorter to free its reservations back to the pool. + drop(sorter); + + // Simulate contention: another partition grabs ALL available + // pool memory. If the merge stream didn't receive the + // pre-reserved bytes via take(), it will fail when it tries + // to allocate memory for reading spill files. + let contender = MemoryConsumer::new("CompetingPartition").register(&pool); + let available = pool_size.saturating_sub(pool.reserved()); + if available > 0 { + contender.try_grow(available).unwrap(); + } + + // The merge stream must still produce correct results despite + // the pool being fully consumed by the contender. This only + // works if sort() transferred the pre-reserved bytes to the + // merge stream (via take()) rather than freeing them. + let batches: Vec<RecordBatch> = merge_stream.try_collect().await?; + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, + (num_batches * 100) as usize, + "Merge stream should produce all rows even under memory contention" + ); + + // Verify data is sorted + let merged = concat_batches(&schema, &batches)?; + let col = merged.column(0).as_primitive::<Int32Type>(); + for i in 1..col.len() { + assert!( + col.value(i - 1) <= col.value(i), + "Output should be sorted, but found {} > {} at index {}", + col.value(i - 1), + col.value(i), + i + ); + } + + drop(contender); + Ok(()) + } + + fn make_sort_exec_with_fetch(fetch: Option<usize>) -> SortExec { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input = Arc::new(EmptyExec::new(schema)); + SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(), + input, + ) + .with_fetch(fetch) + } + + #[test] + fn test_sort_with_fetch_blocks_filter_pushdown() -> Result<()> { + let sort = make_sort_exec_with_fetch(Some(10)); + let desc = sort.gather_filters_for_pushdown( + FilterPushdownPhase::Pre, + vec![Arc::new(Column::new("a", 0))], + &ConfigOptions::new(), + )?; + // Sort with fetch (TopK) must not allow filters to be pushed below it. + assert!(matches!( + desc.parent_filters()[0][0].discriminant, + PushedDown::No + )); + Ok(()) + } + + #[test] + fn test_sort_without_fetch_allows_filter_pushdown() -> Result<()> { + let sort = make_sort_exec_with_fetch(None); + let desc = sort.gather_filters_for_pushdown( + FilterPushdownPhase::Pre, + vec![Arc::new(Column::new("a", 0))], + &ConfigOptions::new(), + )?; + // Plain sort (no fetch) is filter-commutative. + assert!(matches!( + desc.parent_filters()[0][0].discriminant, + PushedDown::Yes + )); + Ok(()) + } + + #[test] + fn test_sort_with_fetch_allows_topk_self_filter_in_post_phase() -> Result<()> { + let sort = make_sort_exec_with_fetch(Some(10)); + assert!(sort.filter.is_some(), "TopK filter should be created"); + + let mut config = ConfigOptions::new(); + config.optimizer.enable_topk_dynamic_filter_pushdown = true; + let desc = sort.gather_filters_for_pushdown( + FilterPushdownPhase::Post, + vec![Arc::new(Column::new("a", 0))], + &config, + )?; + // Parent filters are still blocked in the Post phase. + assert!(matches!( + desc.parent_filters()[0][0].discriminant, + PushedDown::No + )); + // But the TopK self-filter should be pushed down. + assert_eq!(desc.self_filters()[0].len(), 1); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort_preserving_merge.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort_preserving_merge.rs new file mode 100644 index 0000000..ad17f2c --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort_preserving_merge.rs
@@ -0,0 +1,1781 @@ +// 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. + +//! [`SortPreservingMergeExec`] merges multiple sorted streams into one sorted stream. + +use std::sync::Arc; + +use crate::common::spawn_buffered; +use crate::limit::LimitStream; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::projection::{ProjectionExec, make_with_child, update_ordering}; +use crate::sorts::streaming_merge::StreamingMergeBuilder; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, +}; + +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryConsumer; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; + +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; +use log::{debug, trace}; + +/// Sort preserving merge execution plan +/// +/// # Overview +/// +/// This operator implements a K-way merge. It is used to merge multiple sorted +/// streams into a single sorted stream and is highly optimized. +/// +/// ## Inputs: +/// +/// 1. A list of sort expressions +/// 2. An input plan, where each partition is sorted with respect to +/// these sort expressions. +/// +/// ## Output: +/// +/// 1. A single partition that is also sorted with respect to the expressions +/// +/// ## Diagram +/// +/// ```text +/// ┌─────────────────────────┐ +/// │ ┌───┬───┬───┬───┐ │ +/// │ │ A │ B │ C │ D │ ... │──┐ +/// │ └───┴───┴───┴───┘ │ │ +/// └─────────────────────────┘ │ ┌───────────────────┐ ┌───────────────────────────────┐ +/// Stream 1 │ │ │ │ ┌───┬───╦═══╦───┬───╦═══╗ │ +/// ├─▶│SortPreservingMerge│───▶│ │ A │ B ║ B ║ C │ D ║ E ║ ... │ +/// │ │ │ │ └───┴─▲─╩═══╩───┴───╩═══╝ │ +/// ┌─────────────────────────┐ │ └───────────────────┘ └─┬─────┴───────────────────────┘ +/// │ ╔═══╦═══╗ │ │ +/// │ ║ B ║ E ║ ... │──┘ │ +/// │ ╚═══╩═══╝ │ Stable sort if `enable_round_robin_repartition=false`: +/// └─────────────────────────┘ the merged stream places equal rows from stream 1 +/// Stream 2 +/// +/// +/// Input Partitions Output Partition +/// (sorted) (sorted) +/// ``` +/// +/// # Error Handling +/// +/// If any of the input partitions return an error, the error is propagated to +/// the output and inputs are not polled again. +#[derive(Debug, Clone)] +pub struct SortPreservingMergeExec { + /// Input plan with sorted partitions + input: Arc<dyn ExecutionPlan>, + /// Sort expressions + expr: LexOrdering, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Optional number of rows to fetch. Stops producing rows after this fetch + fetch: Option<usize>, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// Use round-robin selection of tied winners of loser tree + /// + /// See [`Self::with_round_robin_repartition`] for more information. + enable_round_robin_repartition: bool, +} + +impl SortPreservingMergeExec { + /// Create a new sort execution plan + pub fn new(expr: LexOrdering, input: Arc<dyn ExecutionPlan>) -> Self { + let cache = Self::compute_properties(&input, expr.clone()); + Self { + input, + expr, + metrics: ExecutionPlanMetricsSet::new(), + fetch: None, + cache: Arc::new(cache), + enable_round_robin_repartition: true, + } + } + + /// Sets the number of rows to fetch + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + /// Sets the selection strategy of tied winners of the loser tree algorithm + /// + /// If true (the default) equal output rows are placed in the merged stream + /// in round robin fashion. This approach consumes input streams at more + /// even rates when there are many rows with the same sort key. + /// + /// If false, equal output rows are always placed in the merged stream in + /// the order of the inputs, resulting in potentially slower execution but a + /// stable output order. + pub fn with_round_robin_repartition( + mut self, + enable_round_robin_repartition: bool, + ) -> Self { + self.enable_round_robin_repartition = enable_round_robin_repartition; + self + } + + /// Input schema + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Sort expressions + pub fn expr(&self) -> &LexOrdering { + &self.expr + } + + /// Fetch + pub fn fetch(&self) -> Option<usize> { + self.fetch + } + + /// Creates the cache object that stores the plan properties + /// such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + ordering: LexOrdering, + ) -> PlanProperties { + let input_partitions = input.output_partitioning().partition_count(); + let (drive, scheduling) = if input_partitions > 1 { + (EvaluationType::Eager, SchedulingType::Cooperative) + } else { + ( + input.properties().evaluation_type, + input.properties().scheduling_type, + ) + }; + + let mut eq_properties = input.equivalence_properties().clone(); + eq_properties.clear_per_partition_constants(); + eq_properties.add_ordering(ordering); + PlanProperties::new( + eq_properties, // Equivalence Properties + Partitioning::UnknownPartitioning(1), // Output Partitioning + input.pipeline_behavior(), // Pipeline Behavior + input.boundedness(), // Boundedness + ) + .with_evaluation_type(drive) + .with_scheduling_type(scheduling) + } +} + +impl DisplayAs for SortPreservingMergeExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "SortPreservingMergeExec: [{}]", self.expr)?; + if let Some(fetch) = self.fetch { + write!(f, ", fetch={fetch}")?; + }; + + Ok(()) + } + DisplayFormatType::TreeRender => { + if let Some(fetch) = self.fetch { + writeln!(f, "limit={fetch}")?; + }; + + for (i, e) in self.expr().iter().enumerate() { + e.fmt_sql(f)?; + if i != self.expr().len() - 1 { + write!(f, ", ")?; + } + } + + Ok(()) + } + } + } +} + +impl ExecutionPlan for SortPreservingMergeExec { + fn name(&self) -> &'static str { + "SortPreservingMergeExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } + + /// Sets the number of rows to fetch + fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + Some(Arc::new(Self { + input: Arc::clone(&self.input), + expr: self.expr.clone(), + metrics: self.metrics.clone(), + fetch: limit, + cache: Arc::clone(&self.cache), + enable_round_robin_repartition: self.enable_round_robin_repartition, + })) + } + + fn with_preserve_order( + &self, + preserve_order: bool, + ) -> Option<Arc<dyn ExecutionPlan>> { + self.input + .with_preserve_order(preserve_order) + .and_then(|new_input| { + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) + .ok() + }) + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false] + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + vec![Some(OrderingRequirements::from(self.expr.clone()))] + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!("Start SortPreservingMergeExec::execute for partition: {partition}"); + assert_eq_or_internal_err!( + partition, + 0, + "SortPreservingMergeExec invalid partition {partition}" + ); + + let input_partitions = self.input.output_partitioning().partition_count(); + trace!( + "Number of input partitions of SortPreservingMergeExec::execute: {input_partitions}" + ); + let schema = self.schema(); + + let reservation = + MemoryConsumer::new(format!("SortPreservingMergeExec[{partition}]")) + .register(&context.runtime_env().memory_pool); + + match input_partitions { + 0 => internal_err!( + "SortPreservingMergeExec requires at least one input partition" + ), + 1 => match self.fetch { + Some(fetch) => { + let stream = self.input.execute(0, context)?; + debug!( + "Done getting stream for SortPreservingMergeExec::execute with 1 input with {fetch}" + ); + Ok(Box::pin(LimitStream::new( + stream, + 0, + Some(fetch), + BaselineMetrics::new(&self.metrics, partition), + ))) + } + None => { + let stream = self.input.execute(0, context); + debug!( + "Done getting stream for SortPreservingMergeExec::execute with 1 input without fetch" + ); + stream + } + }, + _ => { + let receivers = (0..input_partitions) + .map(|partition| { + let stream = + self.input.execute(partition, Arc::clone(&context))?; + Ok(spawn_buffered(stream, 1)) + }) + .collect::<Result<_>>()?; + + debug!( + "Done setting up sender-receiver for SortPreservingMergeExec::execute" + ); + + let result = StreamingMergeBuilder::new() + .with_streams(receivers) + .with_schema(schema) + .with_expressions(&self.expr) + .with_metrics(BaselineMetrics::new(&self.metrics, partition)) + .with_batch_size(context.session_config().batch_size()) + .with_fetch(self.fetch) + .with_reservation(reservation) + .with_round_robin_tie_breaker(self.enable_round_robin_repartition) + .build()?; + + debug!( + "Got stream result from SortPreservingMergeStream::new_from_receivers" + ); + + Ok(result) + } + } + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, _partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(None)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + if self.fetch.is_none() { + CardinalityEffect::Equal + } else { + CardinalityEffect::LowerEqual + } + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + /// Tries to swap the projection with its input [`SortPreservingMergeExec`]. + /// If this is possible, it returns the new [`SortPreservingMergeExec`] whose + /// child is a projection. Otherwise, it returns None. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection does not narrow the schema, we should not try to push it down. + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + + let Some(updated_exprs) = update_ordering(self.expr.clone(), projection.expr())? + else { + return Ok(None); + }; + + Ok(Some(Arc::new( + SortPreservingMergeExec::new( + updated_exprs, + make_with_child(projection, self.input())?, + ) + .with_fetch(self.fetch()), + ))) + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = self + .expr() + .iter() + .map(|e| { + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&e.expr)?)), + asc: !e.options.descending, + nulls_first: e.options.nulls_first, + }), + )), + }) + }) + .collect::<Result<Vec<_>>>()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge( + Box::new(protobuf::SortPreservingMergeExecNode { + input: Some(Box::new(input)), + expr, + fetch: self.fetch().map(|f| f as i64).unwrap_or(-1), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortPreservingMergeExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use arrow::compute::SortOptions; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + let spm = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge, + "SortPreservingMergeExec", + ); + let input = ctx.decode_required_child( + spm.input.as_deref(), + "SortPreservingMergeExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = spm + .expr + .iter() + .map(|e| { + let sort = match &e.expr_type { + Some(protobuf::physical_expr_node::ExprType::Sort(s)) => s, + _ => { + return internal_err!( + "SortPreservingMergeExec expression is not a sort expression" + ); + } + }; + let expr = ctx.decode_required_expr( + sort.expr.as_deref(), + input_schema.as_ref(), + "SortPreservingMergeExec", + "sort expression", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !sort.asc, + nulls_first: sort.nulls_first, + }, + }) + }) + .collect::<Result<Vec<_>>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return internal_err!("SortPreservingMergeExec requires an ordering"); + }; + let fetch = (spm.fetch >= 0).then_some(spm.fetch as usize); + Ok(Arc::new( + SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), + )) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::fmt::Formatter; + use std::pin::Pin; + use std::sync::Mutex; + use std::task::{Context, Poll, Waker, ready}; + use std::time::Duration; + + use super::*; + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::execution_plan::{Boundedness, EmissionType}; + use crate::expressions::col; + use crate::metrics::{MetricValue, Timestamp}; + use crate::repartition::RepartitionExec; + use crate::sorts::sort::SortExec; + use crate::statistics::StatisticsContext; + use crate::stream::RecordBatchReceiverStream; + use crate::test::TestMemoryExec; + use crate::test::exec::{ + BlockingExec, StatisticsExec, assert_strong_count_converges_to_zero, + }; + use crate::test::{self, assert_is_pending, make_partition}; + use crate::{collect, common}; + + use arrow::array::{ + ArrayRef, Int32Array, Int64Array, RecordBatch, StringArray, + TimestampNanosecondArray, + }; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::stats::Precision; + use datafusion_common::test_util::batches_to_string; + use datafusion_common::{ColumnStatistics, assert_batches_eq, exec_err}; + use datafusion_common_runtime::SpawnedTask; + use datafusion_execution::RecordBatchStream; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::EquivalenceProperties; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + + use futures::{FutureExt, Stream, StreamExt}; + use insta::assert_snapshot; + use tokio::time::timeout; + + // The number in the function is highly related to the memory limit we are testing + // any change of the constant should be aware of + fn generate_task_ctx_for_round_robin_tie_breaker( + target_batch_size: usize, + ) -> Result<Arc<TaskContext>> { + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(20_000_000, 1.0) + .build_arc()?; + let mut config = SessionConfig::new(); + config.options_mut().execution.batch_size = + datafusion_common::config::ConfigNonZeroUsize::try_new(target_batch_size)?; + let task_ctx = TaskContext::default() + .with_runtime(runtime) + .with_session_config(config); + Ok(Arc::new(task_ctx)) + } + + // The number in the function is highly related to the memory limit we are testing, + // any change of the constant should be aware of + fn generate_spm_for_round_robin_tie_breaker( + enable_round_robin_repartition: bool, + ) -> Result<Arc<SortPreservingMergeExec>> { + let row_size = 12500; + let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("a"); row_size])); + let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![0; row_size])); + let rb = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)])?; + let schema = rb.schema(); + + let rbs = std::iter::repeat_n(rb, 1024).collect::<Vec<_>>(); + let sort = [ + PhysicalSortExpr { + expr: col("b", &schema)?, + options: Default::default(), + }, + PhysicalSortExpr { + expr: col("c", &schema)?, + options: Default::default(), + }, + ] + .into(); + + let repartition_exec = RepartitionExec::try_new( + TestMemoryExec::try_new_exec(&[rbs], schema, None)?, + Partitioning::RoundRobinBatch(2), + )?; + let spm = SortPreservingMergeExec::new(sort, Arc::new(repartition_exec)) + .with_round_robin_repartition(enable_round_robin_repartition); + Ok(Arc::new(spm)) + } + + #[test] + fn test_fetch_caps_statistics() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(1_000), + total_byte_size: Precision::Exact(8_000), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + schema.clone(), + )); + let sort = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + let spm = SortPreservingMergeExec::new(sort, input).with_fetch(Some(1)); + let statistics = + StatisticsContext::new().compute(&spm, &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Exact(1)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(8)); + assert!(matches!( + spm.cardinality_effect(), + CardinalityEffect::LowerEqual + )); + Ok(()) + } + + #[test] + fn test_no_fetch_preserves_statistics() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input_stats = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Exact(8_000), + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + let input = Arc::new(StatisticsExec::new(input_stats.clone(), schema.clone())); + let sort = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + let spm = SortPreservingMergeExec::new(sort, input); + let statistics = + StatisticsContext::new().compute(&spm, &StatisticsArgs::new())?; + + assert_eq!(*statistics, input_stats); + assert!(matches!(spm.cardinality_effect(), CardinalityEffect::Equal)); + Ok(()) + } + + /// This test verifies that memory usage stays within limits when the tie breaker is enabled. + /// Any errors here could indicate unintended changes in tie breaker logic. + /// + /// Note: If you adjust constants in this test, ensure that memory usage differs + /// based on whether the tie breaker is enabled or disabled. + #[tokio::test(flavor = "multi_thread")] + async fn test_round_robin_tie_breaker_success() -> Result<()> { + let target_batch_size = 12500; + let task_ctx = generate_task_ctx_for_round_robin_tie_breaker(target_batch_size)?; + let spm = generate_spm_for_round_robin_tie_breaker(true)?; + let _collected = collect(spm, task_ctx).await?; + Ok(()) + } + + /// This test verifies that memory usage stays within limits when the tie breaker is enabled. + /// Any errors here could indicate unintended changes in tie breaker logic. + /// + /// Note: If you adjust constants in this test, ensure that memory usage differs + /// based on whether the tie breaker is enabled or disabled. + #[tokio::test(flavor = "multi_thread")] + async fn test_round_robin_tie_breaker_fail() -> Result<()> { + let task_ctx = generate_task_ctx_for_round_robin_tie_breaker(8192)?; + let spm = generate_spm_for_round_robin_tie_breaker(false)?; + let _err = collect(spm, task_ctx).await.unwrap_err(); + Ok(()) + } + + #[tokio::test] + async fn test_merge_interleave() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("c"), + Some("e"), + Some("g"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("b"), + Some("d"), + Some("f"), + Some("h"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2, 6])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + _test_merge( + &[vec![b1], vec![b2]], + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 10 | b | 1970-01-01T00:00:00.000000004 |", + "| 2 | c | 1970-01-01T00:00:00.000000007 |", + "| 20 | d | 1970-01-01T00:00:00.000000006 |", + "| 7 | e | 1970-01-01T00:00:00.000000006 |", + "| 70 | f | 1970-01-01T00:00:00.000000002 |", + "| 9 | g | 1970-01-01T00:00:00.000000005 |", + "| 90 | h | 1970-01-01T00:00:00.000000002 |", + "| 30 | j | 1970-01-01T00:00:00.000000006 |", // input b2 before b1 + "| 3 | j | 1970-01-01T00:00:00.000000008 |", + "+----+---+-------------------------------+", + ], + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_merge_some_overlap() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![70, 90, 30, 100, 110])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("c"), + Some("d"), + Some("e"), + Some("f"), + Some("g"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2, 6])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + _test_merge( + &[vec![b1], vec![b2]], + &[ + "+-----+---+-------------------------------+", + "| a | b | c |", + "+-----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 70 | c | 1970-01-01T00:00:00.000000004 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 90 | d | 1970-01-01T00:00:00.000000006 |", + "| 30 | e | 1970-01-01T00:00:00.000000002 |", + "| 3 | e | 1970-01-01T00:00:00.000000008 |", + "| 100 | f | 1970-01-01T00:00:00.000000002 |", + "| 110 | g | 1970-01-01T00:00:00.000000006 |", + "+-----+---+-------------------------------+", + ], + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_merge_no_overlap() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("f"), + Some("g"), + Some("h"), + Some("i"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2, 6])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + _test_merge( + &[vec![b1], vec![b2]], + &[ + "+----+---+-------------------------------+", + "| a | b | c |", + "+----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 3 | e | 1970-01-01T00:00:00.000000008 |", + "| 10 | f | 1970-01-01T00:00:00.000000004 |", + "| 20 | g | 1970-01-01T00:00:00.000000006 |", + "| 70 | h | 1970-01-01T00:00:00.000000002 |", + "| 90 | i | 1970-01-01T00:00:00.000000002 |", + "| 30 | j | 1970-01-01T00:00:00.000000006 |", + "+----+---+-------------------------------+", + ], + task_ctx, + ) + .await; + } + + #[tokio::test] + async fn test_merge_three_partitions() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("f"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![8, 7, 6, 5, 8])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 70, 90, 30])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("e"), + Some("g"), + Some("h"), + Some("i"), + Some("j"), + ])); + let c: ArrayRef = + Arc::new(TimestampNanosecondArray::from(vec![40, 60, 20, 20, 60])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 700, 900, 300])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + Some("f"), + Some("g"), + Some("h"), + Some("i"), + Some("j"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![4, 6, 2, 2, 6])); + let b3 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + _test_merge( + &[vec![b1], vec![b2], vec![b3]], + &[ + "+-----+---+-------------------------------+", + "| a | b | c |", + "+-----+---+-------------------------------+", + "| 1 | a | 1970-01-01T00:00:00.000000008 |", + "| 2 | b | 1970-01-01T00:00:00.000000007 |", + "| 7 | c | 1970-01-01T00:00:00.000000006 |", + "| 9 | d | 1970-01-01T00:00:00.000000005 |", + "| 10 | e | 1970-01-01T00:00:00.000000040 |", + "| 100 | f | 1970-01-01T00:00:00.000000004 |", + "| 3 | f | 1970-01-01T00:00:00.000000008 |", + "| 200 | g | 1970-01-01T00:00:00.000000006 |", + "| 20 | g | 1970-01-01T00:00:00.000000060 |", + "| 700 | h | 1970-01-01T00:00:00.000000002 |", + "| 70 | h | 1970-01-01T00:00:00.000000020 |", + "| 900 | i | 1970-01-01T00:00:00.000000002 |", + "| 90 | i | 1970-01-01T00:00:00.000000020 |", + "| 300 | j | 1970-01-01T00:00:00.000000006 |", + "| 30 | j | 1970-01-01T00:00:00.000000060 |", + "+-----+---+-------------------------------+", + ], + task_ctx, + ) + .await; + } + + async fn _test_merge( + partitions: &[Vec<RecordBatch>], + exp: &[&str], + context: Arc<TaskContext>, + ) { + let schema = partitions[0][0].schema(); + let sort = [ + PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: Default::default(), + }, + PhysicalSortExpr { + expr: col("c", &schema).unwrap(), + options: Default::default(), + }, + ] + .into(); + let exec = TestMemoryExec::try_new_exec(partitions, schema, None).unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); + + let collected = collect(merge, context).await.unwrap(); + assert_batches_eq!(exp, collected.as_slice()); + } + + async fn sorted_merge( + input: Arc<dyn ExecutionPlan>, + sort: LexOrdering, + context: Arc<TaskContext>, + ) -> RecordBatch { + let merge = Arc::new(SortPreservingMergeExec::new(sort, input)); + let mut result = collect(merge, context).await.unwrap(); + assert_eq!(result.len(), 1); + result.remove(0) + } + + async fn partition_sort( + input: Arc<dyn ExecutionPlan>, + sort: LexOrdering, + context: Arc<TaskContext>, + ) -> RecordBatch { + let sort_exec = + Arc::new(SortExec::new(sort.clone(), input).with_preserve_partitioning(true)); + sorted_merge(sort_exec, sort, context).await + } + + async fn basic_sort( + src: Arc<dyn ExecutionPlan>, + sort: LexOrdering, + context: Arc<TaskContext>, + ) -> RecordBatch { + let merge = Arc::new(CoalescePartitionsExec::new(src)); + let sort_exec = Arc::new(SortExec::new(sort, merge)); + let mut result = collect(sort_exec, context).await.unwrap(); + assert_eq!(result.len(), 1); + result.remove(0) + } + + #[tokio::test] + async fn test_partition_sort() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let partitions = 4; + let csv = test::scan_partitioned(partitions); + let schema = csv.schema(); + + let sort: LexOrdering = [PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions { + descending: true, + nulls_first: true, + }, + }] + .into(); + + let basic = + basic_sort(Arc::clone(&csv), sort.clone(), Arc::clone(&task_ctx)).await; + let partition = partition_sort(csv, sort, Arc::clone(&task_ctx)).await; + + let basic = arrow::util::pretty::pretty_format_batches(&[basic]) + .unwrap() + .to_string(); + let partition = arrow::util::pretty::pretty_format_batches(&[partition]) + .unwrap() + .to_string(); + + assert_eq!( + basic, partition, + "basic:\n\n{basic}\n\npartition:\n\n{partition}\n\n" + ); + + Ok(()) + } + + // Split the provided record batch into multiple batch_size record batches + fn split_batch(sorted: &RecordBatch, batch_size: usize) -> Vec<RecordBatch> { + let batches = sorted.num_rows().div_ceil(batch_size); + + // Split the sorted RecordBatch into multiple + (0..batches) + .map(|batch_idx| { + let columns = (0..sorted.num_columns()) + .map(|column_idx| { + let length = + batch_size.min(sorted.num_rows() - batch_idx * batch_size); + + sorted + .column(column_idx) + .slice(batch_idx * batch_size, length) + }) + .collect(); + + RecordBatch::try_new(sorted.schema(), columns).unwrap() + }) + .collect() + } + + async fn sorted_partitioned_input( + sort: LexOrdering, + sizes: &[usize], + context: Arc<TaskContext>, + ) -> Result<Arc<dyn ExecutionPlan>> { + let partitions = 4; + let csv = test::scan_partitioned(partitions); + + let sorted = basic_sort(csv, sort, context).await; + let split: Vec<_> = sizes.iter().map(|x| split_batch(&sorted, *x)).collect(); + + TestMemoryExec::try_new_exec(&split, sorted.schema(), None).map(|e| e as _) + } + + #[tokio::test] + async fn test_partition_sort_streaming_input() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = make_partition(11).schema(); + let sort: LexOrdering = [PhysicalSortExpr { + expr: col("i", &schema)?, + options: Default::default(), + }] + .into(); + + let input = + sorted_partitioned_input(sort.clone(), &[10, 3, 11], Arc::clone(&task_ctx)) + .await?; + let basic = + basic_sort(Arc::clone(&input), sort.clone(), Arc::clone(&task_ctx)).await; + let partition = sorted_merge(input, sort, Arc::clone(&task_ctx)).await; + + assert_eq!(basic.num_rows(), 1200); + assert_eq!(partition.num_rows(), 1200); + + let basic = arrow::util::pretty::pretty_format_batches(&[basic])?.to_string(); + let partition = + arrow::util::pretty::pretty_format_batches(&[partition])?.to_string(); + + assert_eq!(basic, partition); + + Ok(()) + } + + #[tokio::test] + async fn test_partition_sort_streaming_input_output() -> Result<()> { + let schema = make_partition(11).schema(); + let sort: LexOrdering = [PhysicalSortExpr { + expr: col("i", &schema)?, + options: Default::default(), + }] + .into(); + + // Test streaming with default batch size + let task_ctx = Arc::new(TaskContext::default()); + let input = + sorted_partitioned_input(sort.clone(), &[10, 5, 13], Arc::clone(&task_ctx)) + .await?; + let basic = basic_sort(Arc::clone(&input), sort.clone(), task_ctx).await; + + // batch size of 23 + let task_ctx = TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(23)); + let task_ctx = Arc::new(task_ctx); + + let merge = Arc::new(SortPreservingMergeExec::new(sort, input)); + let merged = collect(merge, task_ctx).await?; + + assert_eq!(merged.len(), 53); + assert_eq!(basic.num_rows(), 1200); + assert_eq!(merged.iter().map(|x| x.num_rows()).sum::<usize>(), 1200); + + let basic = arrow::util::pretty::pretty_format_batches(&[basic])?.to_string(); + let partition = arrow::util::pretty::pretty_format_batches(&merged)?.to_string(); + + assert_eq!(basic, partition); + + Ok(()) + } + + #[tokio::test] + async fn test_nulls() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + None, + Some("a"), + Some("b"), + Some("d"), + Some("e"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![ + Some(8), + None, + Some(6), + None, + Some(4), + ])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![ + None, + Some("b"), + Some("g"), + Some("h"), + Some("i"), + ])); + let c: ArrayRef = Arc::new(TimestampNanosecondArray::from(vec![ + Some(8), + None, + Some(5), + None, + Some(4), + ])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap(); + let schema = b1.schema(); + + let sort = [ + PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + PhysicalSortExpr { + expr: col("c", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: false, + }, + }, + ] + .into(); + let exec = + TestMemoryExec::try_new_exec(&[vec![b1], vec![b2]], schema, None).unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); + + let collected = collect(merge, task_ctx).await.unwrap(); + assert_eq!(collected.len(), 1); + + assert_snapshot!(batches_to_string(collected.as_slice()), @r" + +---+---+-------------------------------+ + | a | b | c | + +---+---+-------------------------------+ + | 1 | | 1970-01-01T00:00:00.000000008 | + | 1 | | 1970-01-01T00:00:00.000000008 | + | 2 | a | | + | 7 | b | 1970-01-01T00:00:00.000000006 | + | 2 | b | | + | 9 | d | | + | 3 | e | 1970-01-01T00:00:00.000000004 | + | 3 | g | 1970-01-01T00:00:00.000000005 | + | 4 | h | | + | 5 | i | 1970-01-01T00:00:00.000000004 | + +---+---+-------------------------------+ + "); + } + + #[tokio::test] + async fn test_sort_merge_single_partition_with_fetch() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])); + let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + let schema = batch.schema(); + + let sort = [PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let exec = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let merge = + Arc::new(SortPreservingMergeExec::new(sort, exec).with_fetch(Some(2))); + + let collected = collect(merge, task_ctx).await.unwrap(); + assert_eq!(collected.len(), 1); + + assert_snapshot!(batches_to_string(collected.as_slice()), @r" + +---+---+ + | a | b | + +---+---+ + | 1 | a | + | 2 | b | + +---+---+ + "); + } + + #[tokio::test] + async fn test_sort_merge_single_partition_without_fetch() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])); + let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + let schema = batch.schema(); + + let sort = [PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let exec = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); + + let collected = collect(merge, task_ctx).await.unwrap(); + assert_eq!(collected.len(), 1); + + assert_snapshot!(batches_to_string(collected.as_slice()), @r" + +---+---+ + | a | b | + +---+---+ + | 1 | a | + | 2 | b | + | 7 | c | + | 9 | d | + | 3 | e | + +---+---+ + "); + } + + #[tokio::test] + async fn test_async() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = make_partition(11).schema(); + let sort: LexOrdering = [PhysicalSortExpr { + expr: col("i", &schema).unwrap(), + options: SortOptions::default(), + }] + .into(); + + let batches = + sorted_partitioned_input(sort.clone(), &[5, 7, 3], Arc::clone(&task_ctx)) + .await?; + + let partition_count = batches.output_partitioning().partition_count(); + let mut streams = Vec::with_capacity(partition_count); + + for partition in 0..partition_count { + let mut builder = RecordBatchReceiverStream::builder(Arc::clone(&schema), 1); + + let sender = builder.tx(); + + let mut stream = batches.execute(partition, Arc::clone(&task_ctx)).unwrap(); + builder.spawn(async move { + while let Some(batch) = stream.next().await { + sender.send(batch).await.unwrap(); + // This causes the MergeStream to wait for more input + tokio::time::sleep(Duration::from_millis(10)).await; + } + + Ok(()) + }); + + streams.push(builder.build()); + } + + let metrics = ExecutionPlanMetricsSet::new(); + let reservation = + MemoryConsumer::new("test").register(&task_ctx.runtime_env().memory_pool); + + let fetch = None; + let merge_stream = StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(batches.schema()) + .with_expressions(&sort) + .with_metrics(BaselineMetrics::new(&metrics, 0)) + .with_batch_size(task_ctx.session_config().batch_size()) + .with_fetch(fetch) + .with_reservation(reservation) + .build()?; + + let mut merged = common::collect(merge_stream).await.unwrap(); + + assert_eq!(merged.len(), 1); + let merged = merged.remove(0); + let basic = basic_sort(batches, sort.clone(), Arc::clone(&task_ctx)).await; + + let basic = arrow::util::pretty::pretty_format_batches(&[basic]) + .unwrap() + .to_string(); + let partition = arrow::util::pretty::pretty_format_batches(&[merged]) + .unwrap() + .to_string(); + + assert_eq!( + basic, partition, + "basic:\n\n{basic}\n\npartition:\n\n{partition}\n\n" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_merge_metrics() { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("a"), Some("c")])); + let b1 = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![10, 20])); + let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("b"), Some("d")])); + let b2 = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + + let schema = b1.schema(); + let sort = [PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: Default::default(), + }] + .into(); + let exec = + TestMemoryExec::try_new_exec(&[vec![b1], vec![b2]], schema, None).unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); + + let collected = collect(Arc::clone(&merge) as Arc<dyn ExecutionPlan>, task_ctx) + .await + .unwrap(); + assert_snapshot!(batches_to_string(collected.as_slice()), @r" + +----+---+ + | a | b | + +----+---+ + | 1 | a | + | 10 | b | + | 2 | c | + | 20 | d | + +----+---+ + "); + + // Now, validate metrics + let metrics = merge.metrics().unwrap(); + + assert_eq!(metrics.output_rows().unwrap(), 4); + assert!(metrics.elapsed_compute().unwrap() > 0); + + let mut saw_start = false; + let mut saw_end = false; + metrics.iter().for_each(|m| match m.value() { + MetricValue::StartTimestamp(ts) => { + saw_start = true; + assert!(nanos_from_timestamp(ts) > 0); + } + MetricValue::EndTimestamp(ts) => { + saw_end = true; + assert!(nanos_from_timestamp(ts) > 0); + } + _ => {} + }); + + assert!(saw_start); + assert!(saw_end); + } + + fn nanos_from_timestamp(ts: &Timestamp) -> i64 { + ts.value().unwrap().timestamp_nanos_opt().unwrap() + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 2)); + let refs = blocking_exec.refs(); + let sort_preserving_merge_exec = Arc::new(SortPreservingMergeExec::new( + [PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions::default(), + }] + .into(), + blocking_exec, + )); + + let fut = collect(sort_preserving_merge_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn test_stable_sort() { + let task_ctx = Arc::new(TaskContext::default()); + + // Create record batches like: + // batch_number |value + // -------------+------ + // 1 | A + // 1 | B + // + // Ensure that the output is in the same order the batches were fed + let partitions: Vec<Vec<RecordBatch>> = (0..10) + .map(|batch_number| { + let batch_number: Int32Array = + vec![Some(batch_number), Some(batch_number)] + .into_iter() + .collect(); + let value: StringArray = vec![Some("A"), Some("B")].into_iter().collect(); + + let batch = RecordBatch::try_from_iter(vec![ + ("batch_number", Arc::new(batch_number) as ArrayRef), + ("value", Arc::new(value) as ArrayRef), + ]) + .unwrap(); + + vec![batch] + }) + .collect(); + + let schema = partitions[0][0].schema(); + + let sort = [PhysicalSortExpr { + expr: col("value", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + + let exec = TestMemoryExec::try_new_exec(&partitions, schema, None).unwrap(); + let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); + + let collected = collect(merge, task_ctx).await.unwrap(); + assert_eq!(collected.len(), 1); + + // Expect the data to be sorted first by "batch_number" (because + // that was the order it was fed in, even though only "value" + // is in the sort key) + assert_snapshot!(batches_to_string(collected.as_slice()), @r" + +--------------+-------+ + | batch_number | value | + +--------------+-------+ + | 0 | A | + | 1 | A | + | 2 | A | + | 3 | A | + | 4 | A | + | 5 | A | + | 6 | A | + | 7 | A | + | 8 | A | + | 9 | A | + | 0 | B | + | 1 | B | + | 2 | B | + | 3 | B | + | 4 | B | + | 5 | B | + | 6 | B | + | 7 | B | + | 8 | B | + | 9 | B | + +--------------+-------+ + "); + } + + #[derive(Debug)] + struct CongestionState { + wakers: Vec<Waker>, + unpolled_partitions: HashSet<usize>, + } + + #[derive(Debug)] + struct Congestion { + congestion_state: Mutex<CongestionState>, + } + + impl Congestion { + fn new(partition_count: usize) -> Self { + Congestion { + congestion_state: Mutex::new(CongestionState { + wakers: vec![], + unpolled_partitions: (0usize..partition_count).collect(), + }), + } + } + + fn check_congested(&self, partition: usize, cx: &mut Context<'_>) -> Poll<()> { + let mut state = self.congestion_state.lock().unwrap(); + + state.unpolled_partitions.remove(&partition); + + if state.unpolled_partitions.is_empty() { + state.wakers.iter().for_each(|w| w.wake_by_ref()); + state.wakers.clear(); + Poll::Ready(()) + } else { + state.wakers.push(cx.waker().clone()); + Poll::Pending + } + } + } + + /// It returns pending for the 2nd partition until the 3rd partition is polled. The 1st + /// partition is exhausted from the start, and if it is polled more than one, it panics. + #[derive(Debug, Clone)] + struct CongestedExec { + schema: Schema, + cache: Arc<PlanProperties>, + congestion: Arc<Congestion>, + } + + impl CongestedExec { + fn compute_properties(schema: SchemaRef) -> PlanProperties { + let columns = schema + .fields + .iter() + .enumerate() + .map(|(i, f)| Arc::new(Column::new(f.name(), i)) as Arc<dyn PhysicalExpr>) + .collect::<Vec<_>>(); + let mut eq_properties = EquivalenceProperties::new(schema); + eq_properties.add_ordering( + columns + .iter() + .map(|expr| PhysicalSortExpr::new_default(Arc::clone(expr))), + ); + PlanProperties::new( + eq_properties, + Partitioning::Hash(columns, 3), + EmissionType::Incremental, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ) + } + } + + impl ExecutionPlan for CongestedExec { + fn name(&self) -> &'static str { + Self::static_name() + } + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(CongestedStream { + schema: Arc::new(self.schema.clone()), + none_polled_once: false, + congestion: Arc::clone(&self.congestion), + partition, + })) + } + } + + impl DisplayAs for CongestedExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "CongestedExec",).unwrap() + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "").unwrap() + } + } + Ok(()) + } + } + + /// It returns pending for the 2nd partition until the 3rd partition is polled. The 1st + /// partition is exhausted from the start, and if it is polled more than once, it panics. + #[derive(Debug)] + pub struct CongestedStream { + schema: SchemaRef, + none_polled_once: bool, + congestion: Arc<Congestion>, + partition: usize, + } + + impl Stream for CongestedStream { + type Item = Result<RecordBatch>; + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + match self.partition { + 0 => { + let _ = self.congestion.check_congested(self.partition, cx); + if self.none_polled_once { + panic!("Exhausted stream is polled more than once") + } else { + self.none_polled_once = true; + Poll::Ready(None) + } + } + _ => { + ready!(self.congestion.check_congested(self.partition, cx)); + Poll::Ready(None) + } + } + } + } + + impl RecordBatchStream for CongestedStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + } + + #[tokio::test] + async fn test_spm_congestion() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); + let properties = CongestedExec::compute_properties(Arc::new(schema.clone())); + let partition_count = properties.output_partitioning().partition_count(); + let source = CongestedExec { + schema: schema.clone(), + cache: Arc::new(properties), + congestion: Arc::new(Congestion::new(partition_count)), + }; + let spm = SortPreservingMergeExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new( + "c1", 0, + )))] + .into(), + Arc::new(source), + ); + let spm_task = SpawnedTask::spawn(collect(Arc::new(spm), task_ctx)); + + let result = timeout(Duration::from_secs(3), spm_task.join()).await; + match result { + Ok(Ok(Ok(_batches))) => Ok(()), + Ok(Ok(Err(e))) => Err(e), + Ok(Err(_)) => exec_err!("SortPreservingMerge task panicked or was cancelled"), + Err(_) => exec_err!("SortPreservingMerge caused a deadlock"), + } + } + + #[tokio::test] + async fn test_sort_merge_stops_after_error_with_buffered_rows() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let sort: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new( + "i", 0, + )) + as Arc<dyn PhysicalExpr>)] + .into(); + + let mut stream0 = RecordBatchReceiverStream::builder(Arc::clone(&schema), 2); + let tx0 = stream0.tx(); + let schema0 = Arc::clone(&schema); + stream0.spawn(async move { + let batch = + RecordBatch::try_new(schema0, vec![Arc::new(Int32Array::from(vec![1]))])?; + tx0.send(Ok(batch)).await.unwrap(); + tx0.send(exec_err!("stream failure")).await.unwrap(); + Ok(()) + }); + + let mut stream1 = RecordBatchReceiverStream::builder(Arc::clone(&schema), 1); + let tx1 = stream1.tx(); + let schema1 = Arc::clone(&schema); + stream1.spawn(async move { + let batch = + RecordBatch::try_new(schema1, vec![Arc::new(Int32Array::from(vec![2]))])?; + tx1.send(Ok(batch)).await.unwrap(); + Ok(()) + }); + + let metrics = ExecutionPlanMetricsSet::new(); + let reservation = + MemoryConsumer::new("test").register(&task_ctx.runtime_env().memory_pool); + + let mut merge_stream = StreamingMergeBuilder::new() + .with_streams(vec![stream0.build(), stream1.build()]) + .with_schema(Arc::clone(&schema)) + .with_expressions(&sort) + .with_metrics(BaselineMetrics::new(&metrics, 0)) + .with_batch_size(task_ctx.session_config().batch_size()) + .with_fetch(None) + .with_reservation(reservation) + .build()?; + + let first = merge_stream.next().await.unwrap(); + assert!(first.is_err(), "expected merge stream to surface the error"); + assert!( + merge_stream.next().await.is_none(), + "merge stream yielded data after returning an error" + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/stream.rs b/native/vendor/datafusion-physical-plan/src/sorts/stream.rs new file mode 100644 index 0000000..1076310 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/stream.rs
@@ -0,0 +1,543 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::sorts::cursor::{ArrayValues, CursorArray, RowValues}; +use crate::{EmptyRecordBatchStream, SendableRecordBatchStream}; +use crate::{PhysicalExpr, PhysicalSortExpr}; +use arrow::array::{Array, UInt32Array}; +use arrow::compute::take_record_batch; +use arrow::datatypes::Schema; +use arrow::record_batch::RecordBatch; +use arrow::row::{RowConverter, Rows, SortField}; +use arrow_ord::sort::lexsort_to_indices; +use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; +use futures::stream::{Fuse, StreamExt}; +use std::iter::FusedIterator; +use std::marker::PhantomData; +use std::mem; +use std::sync::Arc; +use std::task::{Context, Poll, ready}; + +/// A [`Stream`](futures::Stream) that has multiple partitions that can +/// be polled separately but not concurrently +/// +/// Used by sort preserving merge to decouple the cursor merging logic from +/// the source of the cursors, the intention being to allow preserving +/// any row encoding performed for intermediate sorts +pub trait PartitionedStream: std::fmt::Debug + Send { + type Output; + + /// Returns the number of partitions + fn partitions(&self) -> usize; + + fn poll_next( + &mut self, + cx: &mut Context<'_>, + stream_idx: usize, + ) -> Poll<Option<Self::Output>>; +} + +/// A new type wrapper around a set of fused [`SendableRecordBatchStream`] +/// that implements debug, and skips over empty [`RecordBatch`] +struct FusedStreams(Vec<Fuse<SendableRecordBatchStream>>); + +impl std::fmt::Debug for FusedStreams { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("FusedStreams") + .field("num_streams", &self.0.len()) + .finish() + } +} + +impl FusedStreams { + fn poll_next( + &mut self, + cx: &mut Context<'_>, + stream_idx: usize, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + let poll_result = self.0[stream_idx].poll_next_unpin(cx); + match &poll_result { + Poll::Pending => return Poll::Pending, + Poll::Ready(Some(Ok(b))) if b.num_rows() == 0 => continue, + Poll::Ready(Some(Ok(_))) => return poll_result, + Poll::Ready(None) | Poll::Ready(Some(Err(_))) => { + let stream_schema = self.0[stream_idx].get_ref().schema(); + + // Replace the stream with an empty stream, so we can drop memory usage + let empty_stream: SendableRecordBatchStream = + Box::pin(EmptyRecordBatchStream::new(stream_schema)); + self.0[stream_idx] = empty_stream.fuse(); + + return poll_result; + } + } + } + } +} + +/// A pair of `Arc<Rows>` that can be reused +#[derive(Debug)] +struct ReusableRows { + // inner[stream_idx] holds a two Arcs: + // at start of a new poll + // .0 is the rows from the previous poll (at start), + // .1 is the one that is being written to + // at end of a poll, .0 will be swapped with .1, + inner: Vec<[Option<Arc<Rows>>; 2]>, +} + +impl ReusableRows { + // return a Rows for writing, + // does not clone if the existing rows can be reused + fn take_next(&mut self, stream_idx: usize) -> Result<Rows> { + Arc::try_unwrap(self.inner[stream_idx][1].take().unwrap()).map_err(|_| { + internal_datafusion_err!( + "Rows from RowCursorStream is still in use by consumer" + ) + }) + } + // save the Rows + fn save(&mut self, stream_idx: usize, rows: &Arc<Rows>) { + self.inner[stream_idx][1] = Some(Arc::clone(rows)); + // swap the current with the previous one, so that the next poll can reuse the Rows from the previous poll + let [a, b] = &mut self.inner[stream_idx]; + mem::swap(a, b); + } +} + +/// A [`PartitionedStream`] that wraps a set of [`SendableRecordBatchStream`] +/// and computes [`RowValues`] based on the provided [`PhysicalSortExpr`] +/// Note: the stream returns an error if the consumer buffers more than one RowValues (i.e. holds on to two RowValues +/// from the same partition at the same time). +#[derive(Debug)] +pub struct RowCursorStream { + /// Converter to convert output of physical expressions + converter: RowConverter, + /// The physical expressions to sort by + column_expressions: Vec<Arc<dyn PhysicalExpr>>, + /// Input streams + streams: FusedStreams, + /// Tracks the memory used by `converter` + reservation: MemoryReservation, + /// Allocated rows for each partition, we keep two to allow for buffering one + /// in the consumer of the stream + rows: ReusableRows, +} + +impl RowCursorStream { + pub fn try_new( + schema: &Schema, + expressions: &LexOrdering, + streams: Vec<SendableRecordBatchStream>, + reservation: MemoryReservation, + ) -> Result<Self> { + let sort_fields = expressions + .iter() + .map(|expr| { + let data_type = expr.expr.data_type(schema)?; + Ok(SortField::new_with_options(data_type, expr.options)) + }) + .collect::<Result<Vec<_>>>()?; + + let streams: Vec<_> = streams.into_iter().map(|s| s.fuse()).collect(); + let converter = RowConverter::new(sort_fields)?; + let mut rows = Vec::with_capacity(streams.len()); + for _ in &streams { + // Initialize each stream with an empty Rows + rows.push([ + Some(Arc::new(converter.empty_rows(0, 0))), + Some(Arc::new(converter.empty_rows(0, 0))), + ]); + } + Ok(Self { + converter, + reservation, + column_expressions: expressions.iter().map(|x| Arc::clone(&x.expr)).collect(), + streams: FusedStreams(streams), + rows: ReusableRows { inner: rows }, + }) + } + + fn convert_batch( + &mut self, + batch: &RecordBatch, + stream_idx: usize, + ) -> Result<RowValues> { + let cols = evaluate_expressions_to_arrays(&self.column_expressions, batch)?; + + // At this point, ownership should of this Rows should be unique + let mut rows = self.rows.take_next(stream_idx)?; + + rows.clear(); + + self.converter.append(&mut rows, &cols)?; + self.reservation.try_resize(self.converter.size())?; + + let rows = Arc::new(rows); + + self.rows.save(stream_idx, &rows); + + // track the memory in the newly created Rows. + let rows_reservation = self.reservation.new_empty(); + rows_reservation.try_grow(rows.size())?; + Ok(RowValues::new(rows, rows_reservation)) + } +} + +impl PartitionedStream for RowCursorStream { + type Output = Result<(RowValues, RecordBatch)>; + + fn partitions(&self) -> usize { + self.streams.0.len() + } + + fn poll_next( + &mut self, + cx: &mut Context<'_>, + stream_idx: usize, + ) -> Poll<Option<Self::Output>> { + Poll::Ready(ready!(self.streams.poll_next(cx, stream_idx)).map(|r| { + r.and_then(|batch| { + let cursor = self.convert_batch(&batch, stream_idx)?; + Ok((cursor, batch)) + }) + })) + } +} + +/// Specialized stream for sorts on single primitive columns +pub struct FieldCursorStream<T: CursorArray> { + /// The physical expressions to sort by + sort: PhysicalSortExpr, + /// Input streams + streams: FusedStreams, + /// Create new reservations for each array + reservation: MemoryReservation, + phantom: PhantomData<fn(T) -> T>, +} + +impl<T: CursorArray> std::fmt::Debug for FieldCursorStream<T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrimitiveCursorStream") + .field("num_streams", &self.streams) + .finish() + } +} + +impl<T: CursorArray> FieldCursorStream<T> { + pub fn new( + sort: PhysicalSortExpr, + streams: Vec<SendableRecordBatchStream>, + reservation: MemoryReservation, + ) -> Self { + let streams = streams.into_iter().map(|s| s.fuse()).collect(); + Self { + sort, + streams: FusedStreams(streams), + reservation, + phantom: Default::default(), + } + } + + fn convert_batch(&mut self, batch: &RecordBatch) -> Result<ArrayValues<T::Values>> { + let value = self.sort.expr.evaluate(batch)?; + let array = value.into_array(batch.num_rows())?; + let size_in_mem = array.get_buffer_memory_size(); + let array = array.as_any().downcast_ref::<T>().expect("field values"); + let array_reservation = self.reservation.new_empty(); + array_reservation.try_grow(size_in_mem)?; + Ok(ArrayValues::new( + self.sort.options, + array, + array_reservation, + )) + } +} + +impl<T: CursorArray> PartitionedStream for FieldCursorStream<T> { + type Output = Result<(ArrayValues<T::Values>, RecordBatch)>; + + fn partitions(&self) -> usize { + self.streams.0.len() + } + + fn poll_next( + &mut self, + cx: &mut Context<'_>, + stream_idx: usize, + ) -> Poll<Option<Self::Output>> { + Poll::Ready(ready!(self.streams.poll_next(cx, stream_idx)).map(|r| { + r.and_then(|batch| { + let cursor = self.convert_batch(&batch)?; + Ok((cursor, batch)) + }) + })) + } +} + +/// A lazy, memory-efficient sort iterator used as a fallback during aggregate +/// spill when there is not enough memory for an eager sort (which requires ~2x +/// peak memory to hold both the unsorted and sorted copies simultaneously). +/// +/// On the first call to `next()`, a sorted index array (`UInt32Array`) is +/// computed via `lexsort_to_indices`. Subsequent calls yield chunks of +/// `batch_size` rows by `take`-ing from the original batch using slices of +/// this index array. Each `take` copies data for the chunk (not zero-copy), +/// but only one chunk is live at a time since the caller consumes it before +/// requesting the next. Once all rows have been yielded, the original batch +/// and index array are dropped to free memory. +/// +/// The caller must reserve `sizeof(batch) + sizeof(one chunk)` for this iterator, +/// and free the reservation once the iterator is depleted. +pub(crate) struct IncrementalSortIterator { + batch: RecordBatch, + expressions: LexOrdering, + batch_size: usize, + indices: Option<UInt32Array>, + cursor: usize, +} + +impl IncrementalSortIterator { + pub(crate) fn new( + batch: RecordBatch, + expressions: LexOrdering, + batch_size: usize, + ) -> Self { + Self { + batch, + expressions, + batch_size, + cursor: 0, + indices: None, + } + } +} + +impl Iterator for IncrementalSortIterator { + type Item = Result<RecordBatch>; + + fn next(&mut self) -> Option<Self::Item> { + if self.cursor >= self.batch.num_rows() { + return None; + } + + match self.indices.as_ref() { + None => { + let sort_columns = match self + .expressions + .iter() + .map(|expr| expr.evaluate_to_sort_column(&self.batch)) + .collect::<Result<Vec<_>>>() + { + Ok(cols) => cols, + Err(e) => return Some(Err(e)), + }; + + let indices = match lexsort_to_indices(&sort_columns, None) { + Ok(indices) => indices, + Err(e) => return Some(Err(e.into())), + }; + self.indices = Some(indices); + + // Call again, this time it will hit the Some(indices) branch and return the first batch + self.next() + } + Some(indices) => { + let batch_size = self.batch_size.min(self.batch.num_rows() - self.cursor); + + // Perform the take to produce the next batch + let new_batch_indices = indices.slice(self.cursor, batch_size); + let new_batch = match take_record_batch(&self.batch, &new_batch_indices) { + Ok(batch) => batch, + Err(e) => return Some(Err(e.into())), + }; + + self.cursor += batch_size; + + // If this is the last batch, we can release the memory + if self.cursor >= self.batch.num_rows() { + let schema = self.batch.schema(); + let _ = mem::replace(&mut self.batch, RecordBatch::new_empty(schema)); + self.indices = None; + } + + // Return the new batch + Some(Ok(new_batch)) + } + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + let num_rows = self.batch.num_rows(); + let batch_size = self.batch_size; + let num_batches = num_rows.div_ceil(batch_size); + (num_batches, Some(num_batches)) + } +} + +impl FusedIterator for IncrementalSortIterator {} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{AsArray, Int32Array}; + use arrow::datatypes::{DataType, Field, Int32Type}; + use arrow_schema::SchemaRef; + use datafusion_common::DataFusionError; + use datafusion_execution::RecordBatchStream; + use datafusion_physical_expr::expressions::col; + use futures::Stream; + use std::pin::Pin; + + /// Verifies that `take_record_batch` in `IncrementalSortIterator` actually + /// copies the data into a new allocation rather than returning a zero-copy + /// slice of the original batch. If the output arrays were slices, their + /// underlying buffer length would match the original array's length; a true + /// copy will have a buffer sized to fit only the chunk. + #[test] + fn incremental_sort_iterator_copies_data() -> Result<()> { + let original_len = 10; + let batch_size = 3; + + // Build a batch with a single Int32 column of descending values + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a: Int32Array = Int32Array::from(vec![0; original_len]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(col_a)])?; + + // Sort ascending on column "a" + let expressions = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "a", + &batch.schema(), + )?)]) + .unwrap(); + + let mut total_rows = 0; + IncrementalSortIterator::new(batch.clone(), expressions, batch_size).try_for_each( + |result| { + let chunk = result?; + total_rows += chunk.num_rows(); + + // Every output column must be a fresh allocation whose length + // equals the chunk size, NOT the original array length. + chunk.columns().iter().zip(batch.columns()).for_each(|(arr, original_arr)| { + let (_, scalar_buf, _) = arr.as_primitive::<Int32Type>().clone().into_parts(); + let (_, original_scalar_buf, _) = original_arr.as_primitive::<Int32Type>().clone().into_parts(); + + assert_ne!(scalar_buf.inner().data_ptr(), original_scalar_buf.inner().data_ptr(), "Expected a copy of the data for each chunk, but got a slice that shares the same buffer as the original array"); + }); + + Result::<_, DataFusionError>::Ok(()) + }, + )?; + + assert_eq!(total_rows, original_len); + Ok(()) + } + + #[test] + fn test_fused_stream_drop_finished_streams() { + #[derive(Clone)] + struct SingleItemManualStream { + // Held only so its `Arc` strong count reveals when the stream is dropped. + #[expect(dead_code)] + hold_ref: Arc<()>, + record_batch: RecordBatch, + should_finish: bool, + } + + impl Stream for SingleItemManualStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if !self.should_finish { + self.should_finish = true; + return Poll::Ready(Some(Ok(self.record_batch.clone()))); + } + + Poll::Ready(None) + } + } + + impl RecordBatchStream for SingleItemManualStream { + fn schema(&self) -> SchemaRef { + self.record_batch.schema() + } + } + + let hold_ref = Arc::new(()); + let record_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + + let stream_1 = SingleItemManualStream { + hold_ref: Arc::clone(&hold_ref), + should_finish: false, + record_batch: record_batch.clone(), + }; + let stream_2 = stream_1.clone(); + + let stream_1: SendableRecordBatchStream = Box::pin(stream_1); + let stream_2: SendableRecordBatchStream = Box::pin(stream_2); + + let mut fused_stream = FusedStreams(vec![stream_1.fuse(), stream_2.fuse()]); + + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + // The original plus one clone held by each of the two streams. + assert_eq!(Arc::strong_count(&hold_ref), 3); + + // First fetch from stream 0 yields its single batch. + // the stream is not finished yet, so nothing is dropped. + let poll = fused_stream.poll_next(&mut cx, 0); + assert!(matches!(poll, Poll::Ready(Some(Ok(_))))); + assert_eq!(Arc::strong_count(&hold_ref), 3); + + // Second fetch from stream 0 returns `None`, so it is replaced with an + // empty stream and dropped, releasing its `hold_ref` clone. + // running 3 times to make sure the stream is fused correctly + for _ in 0..3 { + let poll = fused_stream.poll_next(&mut cx, 0); + assert!(matches!(poll, Poll::Ready(None))); + assert_eq!(Arc::strong_count(&hold_ref), 2); + } + + // First fetch from stream 1 yields its single batch + // the stream is not finished yet, so nothing is dropped. + let poll = fused_stream.poll_next(&mut cx, 1); + assert!(matches!(poll, Poll::Ready(Some(Ok(_))))); + assert_eq!(Arc::strong_count(&hold_ref), 2); + + // Second fetch from stream 1 returns `None`, so it is replaced with an + // empty stream and dropped, releasing its `hold_ref` clone. + // running 3 times to make sure the stream is fused correctly + for _ in 0..3 { + let poll = fused_stream.poll_next(&mut cx, 1); + assert!(matches!(poll, Poll::Ready(None))); + assert_eq!(Arc::strong_count(&hold_ref), 1); + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/sorts/streaming_merge.rs b/native/vendor/datafusion-physical-plan/src/sorts/streaming_merge.rs new file mode 100644 index 0000000..81adad8 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/sorts/streaming_merge.rs
@@ -0,0 +1,382 @@ +// 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. + +//! Merge that deals with an arbitrary size of streaming inputs. +//! This is an order-preserving merge. + +use crate::metrics::BaselineMetrics; +use crate::sorts::multi_level_merge::MultiLevelMergeBuilder; +use crate::sorts::{ + merge::SortPreservingMergeStream, + stream::{FieldCursorStream, RowCursorStream}, +}; +use crate::{EmptyRecordBatchStream, SendableRecordBatchStream, SpillManager}; +use arrow::array::*; +use arrow::datatypes::{DataType, SchemaRef}; +use datafusion_common::human_readable_size; +use datafusion_common::{Result, assert_or_internal_err, internal_err}; +use datafusion_execution::SpillFile; +use datafusion_execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use std::sync::Arc; + +macro_rules! primitive_merge_helper { + ($t:ty, $($v:ident),+) => { + merge_helper!(PrimitiveArray<$t>, $($v),+) + }; +} + +macro_rules! merge_helper { + ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ + let streams = + FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); + return Ok(SortPreservingMergeStream::new( + Box::new(streams), + $schema, + $tracking_metrics, + $batch_size, + $fetch, + $reservation, + $enable_round_robin_tie_breaker, + ) + .into_stream()); + }}; +} + +pub struct SortedSpillFile { + pub file: Arc<dyn SpillFile>, + + /// how much memory the largest memory batch is taking + pub max_record_batch_memory: usize, +} + +impl std::fmt::Debug for SortedSpillFile { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.file.path() { + Some(path) => write!( + f, + "SortedSpillFile({:?}) takes {}", + path, + human_readable_size(self.max_record_batch_memory) + ), + None => write!( + f, + "SortedSpillFile(<custom_backend>) takes {}", + human_readable_size(self.max_record_batch_memory) + ), + } + } +} + +#[derive(Default)] +pub struct StreamingMergeBuilder<'a> { + streams: Vec<SendableRecordBatchStream>, + sorted_spill_files: Vec<SortedSpillFile>, + spill_manager: Option<SpillManager>, + schema: Option<SchemaRef>, + expressions: Option<&'a LexOrdering>, + metrics: Option<BaselineMetrics>, + batch_size: Option<usize>, + fetch: Option<usize>, + reservation: Option<MemoryReservation>, + enable_round_robin_tie_breaker: bool, +} + +impl<'a> StreamingMergeBuilder<'a> { + pub fn new() -> Self { + Self { + enable_round_robin_tie_breaker: true, + ..Default::default() + } + } + + pub fn with_streams(mut self, streams: Vec<SendableRecordBatchStream>) -> Self { + self.streams = streams; + self + } + + pub fn with_sorted_spill_files( + mut self, + sorted_spill_files: Vec<SortedSpillFile>, + ) -> Self { + self.sorted_spill_files = sorted_spill_files; + self + } + + pub fn with_spill_manager(mut self, spill_manager: SpillManager) -> Self { + self.spill_manager = Some(spill_manager); + self + } + + pub fn with_schema(mut self, schema: SchemaRef) -> Self { + self.schema = Some(schema); + self + } + + pub fn with_expressions(mut self, expressions: &'a LexOrdering) -> Self { + self.expressions = Some(expressions); + self + } + + pub fn with_metrics(mut self, metrics: BaselineMetrics) -> Self { + self.metrics = Some(metrics); + self + } + + pub fn with_batch_size(mut self, batch_size: usize) -> Self { + self.batch_size = Some(batch_size); + self + } + + pub fn with_fetch(mut self, fetch: Option<usize>) -> Self { + self.fetch = fetch; + self + } + + pub fn with_reservation(mut self, reservation: MemoryReservation) -> Self { + self.reservation = Some(reservation); + self + } + + /// See [SortPreservingMergeExec::with_round_robin_repartition] for more + /// information. + /// + /// [SortPreservingMergeExec::with_round_robin_repartition]: crate::sorts::sort_preserving_merge::SortPreservingMergeExec::with_round_robin_repartition + pub fn with_round_robin_tie_breaker( + mut self, + enable_round_robin_tie_breaker: bool, + ) -> Self { + self.enable_round_robin_tie_breaker = enable_round_robin_tie_breaker; + self + } + + /// Bypass the mempool and avoid using the memory reservation. + /// + /// This is not marked as `pub` because it is not recommended to use this method + pub(super) fn with_bypass_mempool(self) -> Self { + let mem_pool: Arc<dyn MemoryPool> = Arc::new(UnboundedMemoryPool::default()); + + self.with_reservation( + MemoryConsumer::new("merge stream mock memory").register(&mem_pool), + ) + } + + pub fn build(self) -> Result<SendableRecordBatchStream> { + let Self { + streams, + sorted_spill_files, + spill_manager, + schema, + metrics, + batch_size, + reservation, + fetch, + expressions, + enable_round_robin_tie_breaker, + } = self; + + // Early return if expressions are empty: + let Some(expressions) = expressions else { + return internal_err!("Sort expressions cannot be empty for streaming merge"); + }; + let schema = schema.expect("Schema cannot be empty for streaming merge"); + + if fetch.is_some_and(|fetch| fetch == 0) { + return Ok(Box::pin(EmptyRecordBatchStream::new(schema))); + } + + let batch_size = + batch_size.expect("Batch size cannot be empty for streaming merge"); + + if batch_size == 0 { + return internal_err!("Batch size cannot be zero for streaming merge"); + } + + if !sorted_spill_files.is_empty() { + // Unwrapping mandatory fields + let metrics = metrics.expect("Metrics cannot be empty for streaming merge"); + let reservation = + reservation.expect("Reservation cannot be empty for streaming merge"); + + return Ok(MultiLevelMergeBuilder::new( + spill_manager.expect("spill_manager should exist"), + schema, + sorted_spill_files, + streams, + expressions.clone(), + metrics, + batch_size, + reservation, + fetch, + enable_round_robin_tie_breaker, + ) + .create_spillable_merge_stream()); + } + + // Early return if streams are empty: + assert_or_internal_err!( + !streams.is_empty(), + "Streams/sorted spill files cannot be empty for streaming merge" + ); + + // Unwrapping mandatory fields + let metrics = metrics.expect("Metrics cannot be empty for streaming merge"); + let reservation = + reservation.expect("Reservation cannot be empty for streaming merge"); + + // Special case single column comparisons with optimized cursor implementations + if expressions.len() == 1 { + let sort = expressions[0].clone(); + let data_type = sort.expr.data_type(schema.as_ref())?; + downcast_primitive! { + data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker), + DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + _ => {} + } + } + + let streams = RowCursorStream::try_new( + schema.as_ref(), + expressions, + streams, + reservation.new_empty(), + )?; + Ok(SortPreservingMergeStream::new( + Box::new(streams), + schema, + metrics, + batch_size, + fetch, + reservation, + enable_round_robin_tie_breaker, + ) + .into_stream()) + } +} + +#[cfg(test)] +mod tests { + use crate::{common::collect, stream::RecordBatchStreamAdapter}; + use std::sync::Arc; + + use super::*; + + use arrow::array::{ArrayRef, RecordBatch}; + use arrow_schema::SortOptions; + use datafusion_common::Result; + use datafusion_execution::TaskContext; + use datafusion_physical_expr::{PhysicalSortExpr, expressions::col}; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_only_1_stream() { + test_fetch_0_should_output_0_rows(1, 0).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_2_streams() { + test_fetch_0_should_output_0_rows(2, 0).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_only_1_spill_file() { + test_fetch_0_should_output_0_rows(0, 1).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_2_spill_files() { + test_fetch_0_should_output_0_rows(0, 2).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_1_stream_and_1_spill_file() { + test_fetch_0_should_output_0_rows(1, 1).await.unwrap(); + } + + async fn test_fetch_0_should_output_0_rows( + number_of_streams: usize, + number_of_spilled_files: usize, + ) -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])); + let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + let schema = batch.schema(); + + let sort: LexOrdering = [PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + + let streams = (0..number_of_streams) + .map(|_| { + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch.clone())]), + )) as SendableRecordBatchStream + }) + .collect::<Vec<SendableRecordBatchStream>>(); + + let spill_manager = SpillManager::new( + task_ctx.runtime_env(), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(&schema), + ); + + let mut sorted_spill_files: Vec<SortedSpillFile> = vec![]; + + for _ in 0..number_of_spilled_files { + let file = spill_manager + .spill_record_batch_and_finish(std::slice::from_ref(&batch), "spill") + .unwrap() + .unwrap(); + sorted_spill_files.push(SortedSpillFile { + file, + max_record_batch_memory: batch.get_array_memory_size(), + }); + } + + let sorted_output_stream = StreamingMergeBuilder::new() + .with_batch_size(100) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + // Just to avoid having to provide memory pool + .with_bypass_mempool() + .with_schema(schema) + .with_streams(streams) + .with_sorted_spill_files(sorted_spill_files) + .with_spill_manager(spill_manager) + .with_expressions(&sort) + // The whole point of the test - fetch is 0 + .with_fetch(Some(0)) + .build() + .unwrap(); + + let collected = collect(sorted_output_stream).await.unwrap(); + let total: usize = collected.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 0, "fetch=Some(0) must emit zero rows, got {total}"); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/spill/in_progress_spill_file.rs b/native/vendor/datafusion-physical-plan/src/spill/in_progress_spill_file.rs new file mode 100644 index 0000000..71d7cce --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/spill/in_progress_spill_file.rs
@@ -0,0 +1,212 @@ +// 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. + +//! Define the `InProgressSpillFile` struct, which represents an in-progress spill file used for writing `RecordBatch`es to disk, created by `SpillManager`. + +use datafusion_common::Result; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use datafusion_common::exec_datafusion_err; +use datafusion_execution::spill_file::SpillFile; + +use super::{ + IPCStreamWriter, gc_view_arrays, + spill_manager::{GetSlicedSize, SpillManager}, +}; + +/// Represents an in-progress spill file used for writing `RecordBatch`es to disk, created by `SpillManager`. +/// Caller is able to use this struct to incrementally append in-memory batches to +/// the file, and then finalize the file by calling the `finish` method. +pub struct InProgressSpillFile { + pub(crate) spill_writer: Arc<SpillManager>, + /// Lazily initialized writer + writer: Option<IPCStreamWriter>, + /// Lazily initialized in-progress file, it will be moved out when the `finish` method is invoked + in_progress_file: Option<Arc<dyn SpillFile>>, +} + +impl InProgressSpillFile { + pub fn new( + spill_writer: Arc<SpillManager>, + in_progress_file: Arc<dyn SpillFile>, + ) -> Self { + Self { + spill_writer, + in_progress_file: Some(in_progress_file), + writer: None, + } + } + + /// Appends a `RecordBatch` to the spill file, initializing the writer if necessary. + /// + /// Before writing, performs GC on StringView/BinaryView arrays to compact backing + /// buffers. When a view array is sliced, it still references the original full buffers, + /// causing massive spill files without GC (see issue #19414: 820MB → 33MB after GC). + /// + /// Returns the post-GC sliced memory size of the batch for memory accounting. + /// + /// # Errors + /// - Returns an error if the file is not active (has been finalized) + /// - Returns an error if appending would exceed the disk usage limit configured + /// by `max_temp_directory_size` in `DiskManager` + pub fn append_batch(&mut self, batch: &RecordBatch) -> Result<usize> { + if self.in_progress_file.is_none() { + return Err(exec_datafusion_err!( + "Append operation failed: No active in-progress file. The file may have already been finalized." + )); + } + + let gc_batch = gc_view_arrays(batch)?; + + if self.writer.is_none() { + // Use the SpillManager's declared schema rather than the batch's schema. + // Individual batches may have different schemas (e.g., different nullability) + // when they come from different branches of a UnionExec. The SpillManager's + // schema represents the canonical schema that all batches should conform to. + let schema = self.spill_writer.schema(); + if let Some(in_progress_file) = &self.in_progress_file { + let spill_writer = in_progress_file.open_writer()?; + + self.writer = Some(IPCStreamWriter::new( + spill_writer, + schema.as_ref(), + self.spill_writer.compression, + )?); + + // Update metrics + self.spill_writer.metrics.spill_file_count.add(1); + let header_bytes = self.writer.as_ref().unwrap().bytes_written(); + self.spill_writer.metrics.spilled_bytes.add(header_bytes); + } + } + if let Some(writer) = &mut self.writer { + // The writer calculates how many serialized bytes were emitted + let (spilled_rows, delta_bytes) = writer.write(&gc_batch)?; + + self.spill_writer.metrics.spilled_rows.add(spilled_rows); + self.spill_writer.metrics.spilled_bytes.add(delta_bytes); + } + gc_batch.get_sliced_size() + } + + pub fn flush(&mut self) -> Result<()> { + if let Some(writer) = &mut self.writer { + writer.flush()?; + } + Ok(()) + } + + /// Returns a reference to the in-progress file, if it exists. + /// This can be used to get the file path for creating readers before the file is finished. + pub fn file(&self) -> Option<&Arc<dyn SpillFile>> { + self.in_progress_file.as_ref() + } + + /// Finalizes the write process, returning the completed `SpillFile`. + /// If there are no batches spilled before, it returns `None`. + pub fn finish(&mut self) -> Result<Option<Arc<dyn SpillFile>>> { + if self.in_progress_file.is_none() && self.writer.is_none() { + return Err(exec_datafusion_err!( + "Finish operation failed: file has already been finalized." + )); + } + if let Some(mut writer) = self.writer.take() { + // Finish the writer and capture any final trailing bytes emitted + let delta_bytes = writer.finish()?; + self.spill_writer.metrics.spilled_bytes.add(delta_bytes); + } else { + return Ok(None); + } + + Ok(self.in_progress_file.take()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + use futures::TryStreamExt; + + #[tokio::test] + async fn test_spill_file_uses_spill_manager_schema() -> Result<()> { + let nullable_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("val", DataType::Int64, true), + ])); + let non_nullable_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("val", DataType::Int64, false), + ])); + + let runtime = Arc::new(RuntimeEnvBuilder::new().build()?); + let metrics_set = ExecutionPlanMetricsSet::new(); + let spill_metrics = SpillMetrics::new(&metrics_set, 0); + let spill_manager = Arc::new(SpillManager::new( + runtime, + spill_metrics, + Arc::clone(&nullable_schema), + )); + + let mut in_progress = spill_manager.create_in_progress_file("test")?; + + // First batch: non-nullable val (simulates literal-0 UNION branch) + let non_nullable_batch = RecordBatch::try_new( + Arc::clone(&non_nullable_schema), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(Int64Array::from(vec![0, 0, 0])), + ], + )?; + in_progress.append_batch(&non_nullable_batch)?; + + // Second batch: nullable val with NULLs (simulates table UNION branch) + let nullable_batch = RecordBatch::try_new( + Arc::clone(&nullable_schema), + vec![ + Arc::new(Int64Array::from(vec![4, 5, 6])), + Arc::new(Int64Array::from(vec![Some(10), None, Some(30)])), + ], + )?; + in_progress.append_batch(&nullable_batch)?; + + let spill_file = in_progress.finish()?.unwrap(); + + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + + // Stream schema should be nullable + assert_eq!(stream.schema(), nullable_schema); + + let batches = stream.try_collect::<Vec<_>>().await?; + assert_eq!(batches.len(), 2); + + // Both batches must have the SpillManager's nullable schema + assert_eq!( + batches[0], + non_nullable_batch.with_schema(Arc::clone(&nullable_schema))? + ); + assert_eq!(batches[1], nullable_batch); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/spill/mod.rs b/native/vendor/datafusion-physical-plan/src/spill/mod.rs new file mode 100644 index 0000000..addcf78 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/spill/mod.rs
@@ -0,0 +1,1527 @@ +// 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. + +//! Defines the spilling functions + +pub(crate) mod in_progress_spill_file; +pub(crate) mod replayable_spill_input; +pub(crate) mod spill_manager; +pub mod spill_pool; +use datafusion_execution::spill_file::SpillWriter; +// Moved for refactor, re-export to keep the public API stable +pub use datafusion_common::utils::memory::get_record_batch_memory_size; +// Re-export SpillManager for doctests only (hidden from public docs) +#[doc(hidden)] +pub use spill_manager::SpillManager; + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::array::{ + Array, ArrayRef, BinaryViewArray, BufferSpec, GenericByteViewArray, StringViewArray, + layout, make_array, +}; +use arrow::buffer::Buffer; +use arrow::datatypes::DataType; +use arrow::datatypes::{ByteViewType, Schema, SchemaRef}; +use arrow::ipc::{ + MetadataVersion, + reader::StreamDecoder, + writer::{IpcWriteOptions, StreamWriter}, +}; +use arrow::record_batch::RecordBatch; +use arrow_data::ArrayDataBuilder; +use arrow_ipc::CompressionType; + +use datafusion_common::Result; +use datafusion_common::config::SpillCompression; +use datafusion_execution::RecordBatchStream; +use datafusion_execution::spill_file::SpillFile; +use futures::Stream; +use log::debug; + +/// Stream that reads spill files from a [`SpillFile`] backend as a stream of [`RecordBatch`]es. +/// Uses [`StreamDecoder`] to decode IPC bytes received from the backend's async byte stream. +/// Backends handle their own threading concerns internally - OS files use +/// `tokio::fs::File` which performs blocking IO per-syscall without holding a thread +/// for the file's lifetime, avoiding deadlocks when concurrent reads exceed thread pool limits. +struct SpillReaderStream { + schema: SchemaRef, + decoder: StreamDecoder, + byte_stream: Pin<Box<dyn Stream<Item = Result<bytes::Bytes>> + Send>>, + is_done: bool, + + /// Maximum memory size observed among spilling sorted record batches. + /// This is used for validation purposes during reading each RecordBatch from spill. + /// For context on why this value is recorded and validated, + /// see `physical_plan/sort/multi_level_merge.rs`. + max_record_batch_memory: Option<usize>, + + /// Holds leftover bytes from a chunk when a batch is yielded early + current_buffer: Buffer, + + /// Keeps the file alive until the stream is dropped + _spill_file: Arc<dyn SpillFile>, + + schema_validated: bool, +} + +// Small margin allowed to accommodate slight memory accounting variation +const SPILL_BATCH_MEMORY_MARGIN: usize = 4096; + +impl SpillReaderStream { + fn new( + schema: SchemaRef, + spill_file: Arc<dyn SpillFile>, + max_record_batch_memory: Option<usize>, + ) -> Result<Self> { + let byte_stream = spill_file.read_stream()?; + // DataFusion controls what it writes so it can trust its own IPC output, + // matching the behavior of the previous StreamReader-based implementation. + let decoder = unsafe { StreamDecoder::new().with_skip_validation(true) }; + Ok(Self { + schema, + decoder, + byte_stream, + max_record_batch_memory, + is_done: false, + current_buffer: Buffer::from(&[]), + _spill_file: spill_file, + schema_validated: false, + }) + } +} + +impl Stream for SpillReaderStream { + type Item = Result<RecordBatch>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let this = self.get_mut(); + + if this.is_done { + return Poll::Ready(None); + } + + loop { + if !this.current_buffer.is_empty() { + match this.decoder.decode(&mut this.current_buffer) { + Ok(Some(batch)) => { + // One-time schema validation on the first decoded batch. + // The IPC stream embeds the writer's schema in its header; + // StreamDecoder surfaces it via the first batch's schema. + // We check here rather than in new() because schema bytes + // only arrive after decoding the IPC header from the stream. + if !this.schema_validated { + this.schema_validated = true; + let actual = batch.schema(); + if actual != this.schema { + this.is_done = true; + return Poll::Ready(Some(Err( + datafusion_common::exec_datafusion_err!( + "Spill file schema mismatch: expected {}, got {}. \ + The caller must use the same SpillManager that created \ + the spill file to read it.", + this.schema, + actual + ), + ))); + } + } + if let Some(max_record_batch_memory) = + this.max_record_batch_memory + { + let actual_size = get_record_batch_memory_size(&batch); + if actual_size + > max_record_batch_memory + SPILL_BATCH_MEMORY_MARGIN + { + debug!( + "Record batch memory usage ({actual_size} bytes) exceeds the expected limit ({max_record_batch_memory} bytes) \n\ + by more than the allowed tolerance ({SPILL_BATCH_MEMORY_MARGIN} bytes).\n\ + This likely indicates a bug in memory accounting during spilling." + ); + } + } + return Poll::Ready(Some(Ok(batch))); + } + Ok(None) => { + // The chunk didn't form a complete message. Arrow consumed the partial bytes + // into its internal scratch pad, leaving our current_buffer completely empty. + // We do nothing and fall through to fetch more data. + } + Err(e) => { + this.is_done = true; + return Poll::Ready(Some(Err(e.into()))); + } + } + } + + match futures::ready!(this.byte_stream.as_mut().poll_next(cx)) { + Some(Ok(chunk)) => { + this.current_buffer = Buffer::from(chunk); + } + Some(Err(e)) => { + this.is_done = true; + return Poll::Ready(Some(Err(e))); + } + None => { + this.is_done = true; + + if let Err(e) = this.decoder.finish() { + return Poll::Ready(Some(Err(e.into()))); + } + return Poll::Ready(None); + } + } + } + } +} + +impl RecordBatchStream for SpillReaderStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// A wrapper that counts the exact compressed IPC bytes written by Arrow. +/// +/// Arrow's `StreamWriter` does not return the number of bytes written during its +/// `write()` calls. To accurately track the `spilled_bytes` metrics (especially +/// when LZ4/ZSTD compression is applied), we must intercept the `std::io::Write` +/// trait boundary to count the final serialized payload size. +pub(crate) struct TrackingSpillWriter { + inner: Box<dyn SpillWriter>, + pub(crate) total_bytes_written: usize, +} + +impl TrackingSpillWriter { + pub fn new(inner: Box<dyn SpillWriter>) -> Self { + Self { + inner, + total_bytes_written: 0, + } + } + + pub fn finish(mut self) -> Result<()> { + self.inner.finish() + } +} + +impl std::io::Write for TrackingSpillWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> { + let n = self.inner.write(buf)?; + + self.total_bytes_written += n; + + Ok(n) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +/// Write in Arrow IPC Stream format to an underlying `SpillWriter` backend. +/// Stream format also supports dictionary replacement. +struct IPCStreamWriter { + /// Inner writer + writer: Option<StreamWriter<TrackingSpillWriter>>, + /// Batches written + num_batches: usize, + /// Rows written + num_rows: usize, + /// Bytes written + num_bytes: usize, +} + +impl IPCStreamWriter { + /// Create new writer + /// + /// # Codec contract + /// + /// `arrow-ipc` must be compiled with the `lz4` and `zstd` features + /// (declared explicitly in `datafusion-physical-plan/Cargo.toml`). If + /// those features are absent, `try_with_compression` will return an + /// error at runtime for [`SpillCompression::Lz4Frame`] and + /// [`SpillCompression::Zstd`] variants. The Cargo dependency keeps this + /// contract local and build-visible during Cargo feature resolution, + /// rather than relying solely on workspace-level feature unification; + /// see #21917. + pub fn new( + spill_writer: Box<dyn SpillWriter>, + schema: &Schema, + spill_compression: SpillCompression, + ) -> Result<Self> { + let metadata_version = MetadataVersion::V5; + // Depending on the schema, some array types such as StringViewArray require larger (16 byte in this case) alignment. + // If the actual buffer layout after IPC read does not satisfy the alignment requirement, + // Arrow ArrayBuilder will copy the buffer into a newly allocated, properly aligned buffer. + // This copying may lead to memory blowup during IPC read due to duplicated buffers. + // To avoid this, we compute the maximum required alignment based on the schema and configure the IPCStreamWriter accordingly. + let alignment = get_max_alignment_for_schema(schema); + let mut write_options = + IpcWriteOptions::try_new(alignment, false, metadata_version)?; + + let compression_type = Option::<CompressionType>::from(spill_compression); + write_options = write_options.try_with_compression(compression_type)?; + + let adapter = TrackingSpillWriter::new(spill_writer); + let writer = StreamWriter::try_new_with_options(adapter, schema, write_options)?; + + Ok(Self { + num_batches: 0, + num_rows: 0, + num_bytes: 0, + writer: Some(writer), + }) + } + + /// Writes a single batch to the IPC stream and updates the internal counters. + /// + /// Returns a tuple containing the change in the number of rows and bytes written. + pub fn write(&mut self, batch: &RecordBatch) -> Result<(usize, usize)> { + let writer = self.writer.as_mut().unwrap(); + + let bytes_before = writer.get_ref().total_bytes_written; + writer.write(batch)?; + let bytes_after = writer.get_ref().total_bytes_written; + self.num_batches += 1; + let delta_num_rows = batch.num_rows(); + self.num_rows += delta_num_rows; + let delta_num_bytes = bytes_after - bytes_before; + self.num_bytes += delta_num_bytes; + Ok((delta_num_rows, delta_num_bytes)) + } + + pub fn flush(&mut self) -> Result<()> { + use std::io::Write; + if let Some(writer) = &mut self.writer { + writer.get_mut().flush()?; + } + Ok(()) + } + + /// Finish the writer. + /// + /// Returns the number of trailing bytes written during the finish operation + /// (e.g., IPC metadata and footers). + pub fn finish(&mut self) -> Result<usize> { + let mut writer = self.writer.take().unwrap(); + + let bytes_before = writer.get_ref().total_bytes_written; + writer.finish()?; // Writes IPC tail + + // Extract the adapter and flush the final bytes + let adapter = writer.into_inner()?; + let bytes_after = adapter.total_bytes_written; + adapter.finish()?; + + Ok(bytes_after - bytes_before) + } + /// Returns the total number of bytes written so far + pub fn bytes_written(&self) -> usize { + self.writer + .as_ref() + .map(|w| w.get_ref().total_bytes_written) + .unwrap_or(0) + } +} + +// Returns the maximum byte alignment required by any field in the schema (>= 8), derived from Arrow buffer layouts. +fn get_max_alignment_for_schema(schema: &Schema) -> usize { + let minimum_alignment = 8; + let mut max_alignment = minimum_alignment; + for field in schema.fields() { + let layout = layout(field.data_type()); + let required_alignment = layout + .buffers + .iter() + .map(|buffer_spec| { + if let BufferSpec::FixedWidth { alignment, .. } = buffer_spec { + *alignment + } else { + minimum_alignment + } + }) + .max() + .unwrap_or(minimum_alignment); + max_alignment = std::cmp::max(max_alignment, required_alignment); + } + max_alignment +} + +/// Size of a single view structure in StringView/BinaryView arrays (in bytes). +/// Each view is 16 bytes: 4 bytes length + 4 bytes prefix + 8 bytes buffer ID/offset. +const VIEW_SIZE_BYTES: usize = 16; + +/// Performs garbage collection on StringView and BinaryView arrays before spilling to reduce memory usage. +/// +/// # Why GC is needed +/// +/// StringView and BinaryView arrays can accumulate significant memory waste when sliced. +/// When a large array is sliced (e.g., taking first 100 rows of 1000), the view array +/// still references the original data buffers containing all 1000 rows of data. +/// +/// For example, in the ClickBench benchmark (issue #19414), repeated slicing of StringView +/// arrays resulted in 820MB of spill files that could be reduced to just 33MB after GC - +/// a 96% reduction in size. +/// +/// # How it works +/// +/// The GC process: +/// 1. Identifies view arrays (StringView/BinaryView) in the batch +/// 2. Checks if their data buffers exceed a memory threshold +/// 3. If exceeded, calls the Arrow `gc()` method which creates new compact buffers +/// containing only the data referenced by the current views +/// 4. Returns a new batch with GC'd arrays (or original arrays if GC not needed) +/// +/// # When GC is triggered +/// +/// GC is only performed when data buffers exceed a threshold (currently 10KB). +/// This balances memory savings against the CPU overhead of garbage collection. +/// Small arrays are passed through unchanged since the GC overhead would exceed +/// any memory savings. +/// +/// # Performance considerations +/// +/// - If no view arrays need compaction, the original batch is cloned cheaply +/// - GC is skipped for small buffers to avoid unnecessary CPU overhead +/// - Nested container types are traversed recursively so view arrays inside +/// `List`, `Map`, `Union`, `Dictionary`, and other child-bearing arrays are compacted too +/// - The Arrow `gc()` method itself is optimized and only copies referenced data +pub(crate) fn gc_view_arrays(batch: &RecordBatch) -> Result<RecordBatch> { + let mut mutated = false; + let mut new_columns: Vec<Arc<dyn Array>> = Vec::with_capacity(batch.num_columns()); + + for array in batch.columns() { + let (gc_array, array_mutated) = gc_array(array)?; + mutated |= array_mutated; + new_columns.push(gc_array); + } + + if mutated { + Ok(RecordBatch::try_new(batch.schema(), new_columns)?) + } else { + Ok(batch.clone()) + } +} + +fn gc_array(array: &ArrayRef) -> Result<(ArrayRef, bool)> { + match array.data_type() { + DataType::Utf8View => { + let string_view = array + .as_any() + .downcast_ref::<StringViewArray>() + .expect("Utf8View array should downcast to StringViewArray"); + if should_gc_view_array(string_view) { + Ok((Arc::new(string_view.gc()) as ArrayRef, true)) + } else { + Ok((Arc::clone(array), false)) + } + } + DataType::BinaryView => { + let binary_view = array + .as_any() + .downcast_ref::<BinaryViewArray>() + .expect("BinaryView array should downcast to BinaryViewArray"); + if should_gc_view_array(binary_view) { + Ok((Arc::new(binary_view.gc()) as ArrayRef, true)) + } else { + Ok((Arc::clone(array), false)) + } + } + _ => gc_array_children(array), + } +} + +fn gc_array_children(array: &ArrayRef) -> Result<(ArrayRef, bool)> { + let data = array.to_data(); + if data.child_data().is_empty() { + return Ok((Arc::clone(array), false)); + } + + let mut mutated = false; + let mut child_data = Vec::with_capacity(data.child_data().len()); + for child in data.child_data() { + let child_array = make_array(child.clone()); + let (gc_child, child_mutated) = gc_array(&child_array)?; + mutated |= child_mutated; + child_data.push(gc_child.to_data()); + } + + if !mutated { + return Ok((Arc::clone(array), false)); + } + + let rebuilt = ArrayDataBuilder::new(data.data_type().clone()) + .len(data.len()) + .offset(data.offset()) + .nulls(data.nulls().cloned()) + .buffers(data.buffers().to_vec()) + .child_data(child_data) + .build()?; + + Ok((make_array(rebuilt), true)) +} + +/// Determines whether a view array should be garbage collected before spilling. +/// +/// Arrow's `gc()` always allocates new compact buffers (it is never a no-op), so we +/// check here to skip the allocation cost when data buffers are small. We subtract +/// the views buffer (16 bytes × n_rows) from `get_buffer_memory_size()` so the +/// threshold tracks non-inline string data rather than row count. +fn should_gc_view_array<T: ByteViewType>(array: &GenericByteViewArray<T>) -> bool { + const MIN_BUFFER_SIZE_FOR_GC: usize = 10 * 1024; // 10KB threshold + + if array.data_buffers().is_empty() { + return false; + } + + let data_buffer_size = array + .get_buffer_memory_size() + .saturating_sub(array.len() * VIEW_SIZE_BYTES); + data_buffer_size > MIN_BUFFER_SIZE_FOR_GC +} + +#[cfg(test)] +fn calculate_string_view_waste_ratio(array: &StringViewArray) -> f64 { + use arrow_data::MAX_INLINE_VIEW_LEN; + calculate_view_waste_ratio(array.len(), array.data_buffers(), |i| { + if !array.is_null(i) { + let value = array.value(i); + if value.len() > MAX_INLINE_VIEW_LEN as usize { + return value.len(); + } + } + 0 + }) +} + +#[cfg(test)] +fn calculate_view_waste_ratio<F>( + len: usize, + data_buffers: &[Buffer], + get_value_size: F, +) -> f64 +where + F: Fn(usize) -> usize, +{ + let total_buffer_size: usize = data_buffers.iter().map(|b| b.capacity()).sum(); + if total_buffer_size == 0 { + return 0.0; + } + + let mut actual_used_size = (0..len).map(get_value_size).sum::<usize>(); + actual_used_size += len * VIEW_SIZE_BYTES; + + let waste = total_buffer_size.saturating_sub(actual_used_size); + waste as f64 / total_buffer_size as f64 +} + +#[cfg(test)] +mod tests { + use super::in_progress_spill_file::InProgressSpillFile; + use super::*; + use crate::common::collect; + use crate::metrics::ExecutionPlanMetricsSet; + use crate::metrics::SpillMetrics; + use crate::spill::spill_manager::SpillManager; + use crate::test::build_table_i32; + use arrow::array::{ArrayRef, Int32Array, StringArray}; + use arrow::compute::cast; + use arrow::datatypes::{DataType, Field}; + use datafusion_execution::runtime_env::RuntimeEnv; + use futures::StreamExt as _; + + #[tokio::test] + async fn test_batch_spill_and_read() -> Result<()> { + let batch1 = build_table_i32( + ("a2", &vec![0, 1, 2]), + ("b2", &vec![3, 4, 5]), + ("c2", &vec![4, 5, 6]), + ); + + let batch2 = build_table_i32( + ("a2", &vec![10, 11, 12]), + ("b2", &vec![13, 14, 15]), + ("c2", &vec![14, 15, 16]), + ); + + let schema = batch1.schema(); + let num_rows = batch1.num_rows() + batch2.num_rows(); + + // Construct SpillManager + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema)); + + let spill_file = spill_manager + .spill_record_batch_and_finish(&[batch1, batch2], "Test")? + .unwrap(); + assert!(spill_file.path().unwrap().exists()); + let spilled_rows = spill_manager.metrics.spilled_rows.value(); + assert_eq!(spilled_rows, num_rows); + + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + assert_eq!(stream.schema(), schema); + + let batches = collect(stream).await?; + assert_eq!(batches.len(), 2); + + Ok(()) + } + + #[tokio::test] + async fn test_batch_spill_and_read_dictionary_arrays() -> Result<()> { + // See https://github.com/apache/datafusion/issues/4658 + + let batch1 = build_table_i32( + ("a2", &vec![0, 1, 2]), + ("b2", &vec![3, 4, 5]), + ("c2", &vec![4, 5, 6]), + ); + + let batch2 = build_table_i32( + ("a2", &vec![10, 11, 12]), + ("b2", &vec![13, 14, 15]), + ("c2", &vec![14, 15, 16]), + ); + + // Dictionary encode the arrays + let dict_type = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32)); + let dict_schema = Arc::new(Schema::new(vec![ + Field::new("a2", dict_type.clone(), true), + Field::new("b2", dict_type.clone(), true), + Field::new("c2", dict_type.clone(), true), + ])); + + let batch1 = RecordBatch::try_new( + Arc::clone(&dict_schema), + batch1 + .columns() + .iter() + .map(|array| cast(array, &dict_type)) + .collect::<Result<_, _>>()?, + )?; + + let batch2 = RecordBatch::try_new( + Arc::clone(&dict_schema), + batch2 + .columns() + .iter() + .map(|array| cast(array, &dict_type)) + .collect::<Result<_, _>>()?, + )?; + + // Construct SpillManager + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&dict_schema)); + + let num_rows = batch1.num_rows() + batch2.num_rows(); + let spill_file = spill_manager + .spill_record_batch_and_finish(&[batch1, batch2], "Test")? + .unwrap(); + let spilled_rows = spill_manager.metrics.spilled_rows.value(); + assert_eq!(spilled_rows, num_rows); + + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + assert_eq!(stream.schema(), dict_schema); + let batches = collect(stream).await?; + assert_eq!(batches.len(), 2); + + Ok(()) + } + + #[tokio::test] + async fn test_batch_spill_by_size() -> Result<()> { + let batch1 = build_table_i32( + ("a2", &vec![0, 1, 2, 3]), + ("b2", &vec![3, 4, 5, 6]), + ("c2", &vec![4, 5, 6, 7]), + ); + + let schema = batch1.schema(); + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema)); + + let row_batches: Vec<RecordBatch> = + (0..batch1.num_rows()).map(|i| batch1.slice(i, 1)).collect(); + let (spill_file, max_batch_mem) = spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + row_batches.iter().map(Ok), + "Test Spill", + )? + .unwrap(); + assert!(spill_file.path().unwrap().exists()); + assert!(max_batch_mem > 0); + + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + assert_eq!(stream.schema(), schema); + + let batches = collect(stream).await?; + assert_eq!(batches.len(), 4); + + Ok(()) + } + + fn build_compressible_batch() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, true), + ])); + + let a: ArrayRef = Arc::new(StringArray::from_iter_values(std::iter::repeat_n( + "repeated", 100, + ))); + let b: ArrayRef = Arc::new(Int32Array::from(vec![1; 100])); + let c: ArrayRef = Arc::new(Int32Array::from(vec![2; 100])); + + RecordBatch::try_new(schema, vec![a, b, c]).unwrap() + } + + async fn validate( + spill_manager: &SpillManager, + spill_file: Arc<dyn SpillFile>, + num_rows: usize, + schema: SchemaRef, + batch_count: usize, + ) -> Result<()> { + let spilled_rows = spill_manager.metrics.spilled_rows.value(); + assert_eq!(spilled_rows, num_rows); + + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + assert_eq!(stream.schema(), schema); + + let batches = collect(stream).await?; + assert_eq!(batches.len(), batch_count); + + Ok(()) + } + + #[tokio::test] + async fn test_spill_compression() -> Result<()> { + let batch = build_compressible_batch(); + let num_rows = batch.num_rows(); + let schema = batch.schema(); + let batch_count = 1; + let batches = [batch]; + + // Construct SpillManager + let env = Arc::new(RuntimeEnv::default()); + let uncompressed_metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let lz4_metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let zstd_metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let uncompressed_spill_manager = SpillManager::new( + Arc::clone(&env), + uncompressed_metrics, + Arc::clone(&schema), + ); + let lz4_spill_manager = + SpillManager::new(Arc::clone(&env), lz4_metrics, Arc::clone(&schema)) + .with_compression_type(SpillCompression::Lz4Frame); + let zstd_spill_manager = + SpillManager::new(env, zstd_metrics, Arc::clone(&schema)) + .with_compression_type(SpillCompression::Zstd); + let uncompressed_spill_file = uncompressed_spill_manager + .spill_record_batch_and_finish(&batches, "Test")? + .unwrap(); + let lz4_spill_file = lz4_spill_manager + .spill_record_batch_and_finish(&batches, "Lz4_Test")? + .unwrap(); + let zstd_spill_file = zstd_spill_manager + .spill_record_batch_and_finish(&batches, "ZSTD_Test")? + .unwrap(); + assert!(uncompressed_spill_file.path().unwrap().exists()); + assert!(lz4_spill_file.path().unwrap().exists()); + assert!(zstd_spill_file.path().unwrap().exists()); + + let lz4_spill_size = std::fs::metadata(lz4_spill_file.path().unwrap())?.len(); + let zstd_spill_size = std::fs::metadata(zstd_spill_file.path().unwrap())?.len(); + let uncompressed_spill_size = + std::fs::metadata(uncompressed_spill_file.path().unwrap())?.len(); + + assert!(uncompressed_spill_size > lz4_spill_size); + assert!(uncompressed_spill_size > zstd_spill_size); + + validate( + &lz4_spill_manager, + lz4_spill_file, + num_rows, + Arc::clone(&schema), + batch_count, + ) + .await?; + validate( + &zstd_spill_manager, + zstd_spill_file, + num_rows, + Arc::clone(&schema), + batch_count, + ) + .await?; + validate( + &uncompressed_spill_manager, + uncompressed_spill_file, + num_rows, + schema, + batch_count, + ) + .await?; + Ok(()) + } + + // ==== Spill manager tests ==== + + #[test] + fn test_spill_manager_spill_record_batch_and_finish() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema)); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + )?; + + let temp_file = spill_manager.spill_record_batch_and_finish(&[batch], "Test")?; + assert!(temp_file.is_some()); + assert!(temp_file.unwrap().path().unwrap().exists()); + Ok(()) + } + + fn verify_metrics( + in_progress_file: &InProgressSpillFile, + expected_spill_file_count: usize, + expected_spilled_bytes: usize, + expected_spilled_rows: usize, + ) -> Result<()> { + let actual_spill_file_count = in_progress_file + .spill_writer + .metrics + .spill_file_count + .value(); + let actual_spilled_bytes = + in_progress_file.spill_writer.metrics.spilled_bytes.value(); + let actual_spilled_rows = + in_progress_file.spill_writer.metrics.spilled_rows.value(); + + assert_eq!( + actual_spill_file_count, expected_spill_file_count, + "Spill file count mismatch" + ); + assert_eq!( + actual_spilled_bytes, expected_spilled_bytes, + "Spilled bytes mismatch" + ); + assert_eq!( + actual_spilled_rows, expected_spilled_rows, + "Spilled rows mismatch" + ); + + Ok(()) + } + + #[test] + fn test_in_progress_spill_file_append_and_finish() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let spill_manager = + Arc::new(SpillManager::new(env, metrics, Arc::clone(&schema))); + let mut in_progress_file = spill_manager.create_in_progress_file("Test")?; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + )?; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![4, 5, 6])), + Arc::new(StringArray::from(vec!["d", "e", "f"])), + ], + )?; + // After appending each batch, spilled_rows and spilled_bytes should increase incrementally, + // while spill_file_count remains 1 (since we're writing to the same file) + in_progress_file.append_batch(&batch1)?; + verify_metrics(&in_progress_file, 1, 440, 3)?; + + in_progress_file.append_batch(&batch2)?; + verify_metrics(&in_progress_file, 1, 704, 6)?; + + let completed_file = in_progress_file.finish()?; + assert!(completed_file.is_some()); + assert!(completed_file.unwrap().path().unwrap().exists()); + verify_metrics(&in_progress_file, 1, 712, 6)?; + // Double finish produce error + let result = in_progress_file.finish(); + assert!(result.is_err()); + + Ok(()) + } + + // Test write no batches + #[test] + fn test_in_progress_spill_file_write_no_batches() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let spill_manager = + Arc::new(SpillManager::new(env, metrics, Arc::clone(&schema))); + + // Test write empty batch with interface `InProgressSpillFile` and `append_batch()` + let mut in_progress_file = spill_manager.create_in_progress_file("Test")?; + let completed_file = in_progress_file.finish()?; + assert!(completed_file.is_none()); + + // Test write empty batch with interface `spill_record_batch_and_finish()` + let completed_file = spill_manager.spill_record_batch_and_finish(&[], "Test")?; + assert!(completed_file.is_none()); + + // Test write empty batch with interface `spill_record_batch_iter_and_return_max_batch_memory()` + let empty_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(Vec::<Option<i32>>::new())), + Arc::new(StringArray::from(Vec::<Option<&str>>::new())), + ], + )?; + let completed_file = spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + std::iter::once(Ok(&empty_batch)), + "Test", + )?; + assert!(completed_file.is_none()); + + Ok(()) + } + + #[test] + fn test_reading_more_spills_than_tokio_blocking_threads() -> Result<()> { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap() + .block_on(async { + let batch = build_table_i32( + ("a2", &vec![0, 1, 2]), + ("b2", &vec![3, 4, 5]), + ("c2", &vec![4, 5, 6]), + ); + + let schema = batch.schema(); + + // Construct SpillManager + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema)); + let batches: [_; 10] = std::array::from_fn(|_| batch.clone()); + + let spill_file_1 = spill_manager + .spill_record_batch_and_finish(&batches, "Test1")? + .unwrap(); + let spill_file_2 = spill_manager + .spill_record_batch_and_finish(&batches, "Test2")? + .unwrap(); + + let mut stream_1 = + spill_manager.read_spill_as_stream(spill_file_1, None)?; + let mut stream_2 = + spill_manager.read_spill_as_stream(spill_file_2, None)?; + stream_1.next().await; + stream_2.next().await; + + Ok(()) + }) + } + + #[test] + fn test_alignment_for_schema() -> Result<()> { + let schema = Schema::new(vec![Field::new("strings", DataType::Utf8View, false)]); + let alignment = get_max_alignment_for_schema(&schema); + assert_eq!(alignment, 16); + + let schema = Schema::new(vec![ + Field::new("int32", DataType::Int32, false), + Field::new("int64", DataType::Int64, false), + ]); + let alignment = get_max_alignment_for_schema(&schema); + assert_eq!(alignment, 8); + Ok(()) + } + #[tokio::test] + async fn test_real_time_spill_metrics() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let spill_manager = Arc::new(SpillManager::new( + Arc::clone(&env), + metrics.clone(), + Arc::clone(&schema), + )); + let mut in_progress_file = spill_manager.create_in_progress_file("Test")?; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + )?; + + // Before any batch, metrics should be 0 + assert_eq!(metrics.spilled_bytes.value(), 0); + assert_eq!(metrics.spill_file_count.value(), 0); + + // Append first batch + in_progress_file.append_batch(&batch1)?; + + // Metrics should be updated immediately (at least schema and first batch) + let bytes_after_batch1 = metrics.spilled_bytes.value(); + assert_eq!(bytes_after_batch1, 440); + assert_eq!(metrics.spill_file_count.value(), 1); + + // Check global progress + let progress = env.spilling_progress(); + assert_eq!(progress.current_bytes, bytes_after_batch1 as u64); + assert_eq!(progress.active_files_count, 1); + + // Append another batch + in_progress_file.append_batch(&batch1)?; + let bytes_after_batch2 = metrics.spilled_bytes.value(); + assert!(bytes_after_batch2 > bytes_after_batch1); + + // Check global progress again + let progress = env.spilling_progress(); + assert_eq!(progress.current_bytes, bytes_after_batch2 as u64); + + // Finish the file + let spilled_file = in_progress_file.finish()?; + let final_bytes = metrics.spilled_bytes.value(); + assert!(final_bytes > bytes_after_batch2); + + // Even after finish, file is still "active" until dropped + let progress = env.spilling_progress(); + assert!(progress.current_bytes > 0); + assert_eq!(progress.active_files_count, 1); + + drop(spilled_file); + assert_eq!(env.spilling_progress().active_files_count, 0); + assert_eq!(env.spilling_progress().current_bytes, 0); + + Ok(()) + } + + #[test] + fn test_gc_string_view_before_spill() -> Result<()> { + use arrow::array::StringViewArray; + + let strings: Vec<String> = (0..200) + .map(|i| { + if i % 2 == 0 { + "short_string".to_string() + } else { + "this_is_a_much_longer_string_that_will_not_be_inlined".to_string() + } + }) + .collect(); + + let string_array = StringViewArray::from(strings); + let schema = Arc::new(Schema::new(vec![Field::new( + "strings", + DataType::Utf8View, + false, + )])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(string_array) as ArrayRef], + )?; + let sliced_batch = batch.slice(0, 20); + let gc_batch = gc_view_arrays(&sliced_batch)?; + + assert_eq!(gc_batch.num_rows(), sliced_batch.num_rows()); + assert_eq!(gc_batch.num_columns(), sliced_batch.num_columns()); + + Ok(()) + } + + #[test] + fn test_gc_binary_view_before_spill() -> Result<()> { + use arrow::array::BinaryViewArray; + + let binaries: Vec<Vec<u8>> = (0..200) + .map(|i| { + if i % 2 == 0 { + vec![1, 2, 3, 4] + } else { + vec![1; 50] + } + }) + .collect(); + + let binary_array = + BinaryViewArray::from_iter(binaries.iter().map(|b| Some(b.as_slice()))); + let schema = Arc::new(Schema::new(vec![Field::new( + "binaries", + DataType::BinaryView, + false, + )])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(binary_array) as ArrayRef], + )?; + let sliced_batch = batch.slice(0, 20); + let gc_batch = gc_view_arrays(&sliced_batch)?; + + assert_eq!(gc_batch.num_rows(), sliced_batch.num_rows()); + assert_eq!(gc_batch.num_columns(), sliced_batch.num_columns()); + + Ok(()) + } + + #[test] + fn test_gc_skips_small_arrays() -> Result<()> { + use arrow::array::StringViewArray; + + let strings: Vec<String> = (0..10).map(|i| format!("string_{i}")).collect(); + + let string_array = StringViewArray::from(strings); + let array_ref: ArrayRef = Arc::new(string_array); + + let schema = Arc::new(Schema::new(vec![Field::new( + "strings", + DataType::Utf8View, + false, + )])); + + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array_ref])?; + + // GC should return the original batch for small arrays + let should_gc = should_gc_view_array( + batch + .column(0) + .as_any() + .downcast_ref::<StringViewArray>() + .unwrap(), + ); + let gc_batch = gc_view_arrays(&batch)?; + + assert!(!should_gc); + assert_eq!(gc_batch.num_rows(), batch.num_rows()); + assert!(Arc::ptr_eq(batch.column(0), gc_batch.column(0))); + + Ok(()) + } + + #[test] + fn test_gc_with_mixed_columns() -> Result<()> { + use arrow::array::{Int32Array, StringViewArray}; + + let strings: Vec<String> = (0..200) + .map(|i| format!("long_string_for_gc_testing_{i}")) + .collect(); + + let string_array = StringViewArray::from(strings); + let int_array = Int32Array::from((0..200).collect::<Vec<i32>>()); + + let schema = Arc::new(Schema::new(vec![ + Field::new("strings", DataType::Utf8View, false), + Field::new("ints", DataType::Int32, false), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(string_array) as ArrayRef, + Arc::new(int_array) as ArrayRef, + ], + )?; + + let sliced_batch = batch.slice(0, 50); + let gc_batch = gc_view_arrays(&sliced_batch)?; + + assert_eq!(gc_batch.num_columns(), 2); + assert_eq!(gc_batch.num_rows(), 50); + + Ok(()) + } + + #[test] + fn test_verify_gc_triggers_for_sliced_arrays() -> Result<()> { + let strings: Vec<String> = (0..200) + .map(|i| { + format!( + "http://example.com/very/long/path/that/exceeds/inline/threshold/{i}" + ) + }) + .collect(); + + let string_array = StringViewArray::from(strings); + let schema = Arc::new(Schema::new(vec![Field::new( + "url", + DataType::Utf8View, + false, + )])); + + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(string_array.clone()) as ArrayRef], + )?; + + let sliced = batch.slice(0, 20); + + let sliced_array = sliced + .column(0) + .as_any() + .downcast_ref::<StringViewArray>() + .unwrap(); + let should_gc = should_gc_view_array(sliced_array); + let waste_ratio = calculate_string_view_waste_ratio(sliced_array); + + assert!( + waste_ratio > 0.8, + "Waste ratio should be > 0.8 for sliced array" + ); + assert!( + should_gc, + "GC should trigger for sliced array with high waste" + ); + + Ok(()) + } + + #[test] + fn test_reproduce_issue_19414_string_view_spill_without_gc() -> Result<()> { + use arrow::array::StringViewArray; + use std::fs; + + let num_rows = 1000; + let mut strings = Vec::with_capacity(num_rows); + + for i in 0..num_rows { + let url = match i % 5 { + 0 => format!( + "http://irr.ru/index.php?showalbum/login-leniya7777294,938303130/{i}" + ), + 1 => format!("http://komme%2F27.0.1453.116/very/long/path/{i}"), + 2 => format!("https://produkty%2Fproduct/category/item/{i}"), + 3 => format!( + "http://irr.ru/index.php?showalbum/login-kapusta-advert2668/{i}" + ), + 4 => format!( + "http://irr.ru/index.php?showalbum/login-kapustic/product/{i}" + ), + _ => unreachable!(), + }; + strings.push(url); + } + + let string_array = StringViewArray::from(strings); + let schema = Arc::new(Schema::new(vec![Field::new( + "URL", + DataType::Utf8View, + false, + )])); + + let original_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(string_array.clone()) as ArrayRef], + )?; + + let total_buffer_size: usize = string_array + .data_buffers() + .iter() + .map(|buffer| buffer.capacity()) + .sum(); + + let mut sliced_batches = Vec::new(); + let slice_size = 100; + + for i in (0..num_rows).step_by(slice_size) { + let len = std::cmp::min(slice_size, num_rows - i); + let sliced = original_batch.slice(i, len); + sliced_batches.push(sliced); + } + + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, schema); + + let mut in_progress_file = spill_manager.create_in_progress_file("Test GC")?; + + for batch in &sliced_batches { + in_progress_file.append_batch(batch)?; + } + + let spill_file = in_progress_file.finish()?.unwrap(); + let file_size = fs::metadata(spill_file.path().unwrap())?.len() as usize; + + let theoretical_without_gc = total_buffer_size * sliced_batches.len(); + let reduction_percent = ((theoretical_without_gc - file_size) as f64 + / theoretical_without_gc as f64) + * 100.0; + + assert!( + reduction_percent > 80.0, + "GC should reduce spill file size by >80%, got {reduction_percent:.1}%" + ); + + Ok(()) + } + + #[test] + fn test_spill_with_and_without_gc_comparison() -> Result<()> { + let num_rows = 400; + let strings: Vec<String> = (0..num_rows) + .map(|i| { + format!( + "http://example.com/this/is/a/long/url/path/that/wont/be/inlined/{i}" + ) + }) + .collect(); + + let string_array = StringViewArray::from(strings); + let schema = Arc::new(Schema::new(vec![Field::new( + "url", + DataType::Utf8View, + false, + )])); + + let batch = + RecordBatch::try_new(schema, vec![Arc::new(string_array) as ArrayRef])?; + + let sliced_batch = batch.slice(0, 40); + + let array_without_gc = sliced_batch + .column(0) + .as_any() + .downcast_ref::<StringViewArray>() + .unwrap(); + let size_without_gc: usize = array_without_gc + .data_buffers() + .iter() + .map(|buffer| buffer.capacity()) + .sum(); + + let gc_batch = gc_view_arrays(&sliced_batch)?; + let array_with_gc = gc_batch + .column(0) + .as_any() + .downcast_ref::<StringViewArray>() + .unwrap(); + let size_with_gc: usize = array_with_gc + .data_buffers() + .iter() + .map(|buffer| buffer.capacity()) + .sum(); + + let reduction_percent = + ((size_without_gc - size_with_gc) as f64 / size_without_gc as f64) * 100.0; + + assert!( + reduction_percent > 85.0, + "Expected >85% reduction for 10% slice, got {reduction_percent:.1}%" + ); + + Ok(()) + } + + #[test] + fn test_gc_recurses_into_nested_view_arrays() -> Result<()> { + use arrow::array::{DictionaryArray, Int32Array}; + use arrow::buffer::Buffer; + + let strings: Vec<String> = (0..200) + .map(|i| format!("http://example.com/nested/path/that/is/not/inlined/{i}")) + .collect(); + let string_values = Arc::new(StringViewArray::from(strings)) as ArrayRef; + + let list_data = ArrayDataBuilder::new(DataType::List(Arc::new( + Field::new_list_field(DataType::Utf8View, true), + ))) + .len(20) + .buffers(vec![Buffer::from_iter((0..=20).map(|i| i * 5_i32))]) + .child_data(vec![string_values.slice(0, 100).to_data()]) + .build()?; + let list_array = make_array(list_data); + + let keys = Int32Array::from_iter_values(0..20); + let dictionary = DictionaryArray::new(keys, string_values.slice(0, 20)); + let dictionary_array = Arc::new(dictionary) as ArrayRef; + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "list_strings", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8View, true))), + false, + ), + Field::new( + "dictionary_strings", + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8View), + ), + false, + ), + ])); + let batch = RecordBatch::try_new(schema, vec![list_array, dictionary_array])?; + let gc_batch = gc_view_arrays(&batch)?; + + let gc_list_values = gc_batch.column(0).to_data().child_data()[0].clone(); + let gc_list_values = make_array(gc_list_values); + let gc_list_values = gc_list_values + .as_any() + .downcast_ref::<StringViewArray>() + .unwrap(); + assert!( + calculate_string_view_waste_ratio(gc_list_values) < 0.2, + "GC should compact nested List child views" + ); + + let gc_dictionary_values = gc_batch.column(1).to_data().child_data()[0].clone(); + let gc_dictionary_values = make_array(gc_dictionary_values); + let gc_dictionary_values = gc_dictionary_values + .as_any() + .downcast_ref::<StringViewArray>() + .unwrap(); + assert!( + calculate_string_view_waste_ratio(gc_dictionary_values) < 0.2, + "GC should compact nested Dictionary values" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_spill_file_size_gc_verification_string_view() -> Result<()> { + use arrow::array::StringViewArray; + use std::fs; + + // 1. Setup bloated data (large buffers) + let num_rows = 1000; + let string_array: StringViewArray = (0..num_rows) + .map(|i| Some(format!("this_is_a_long_string_to_ensure_it_is_not_inlined_and_causes_waste_{i}"))) + .collect(); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(string_array.clone()) as ArrayRef], + )?; + + // 2. Slice it heavily (1% of the data) + let sliced_batch = batch.slice(0, 10); + + // 3. Spill to disk using SpillManager + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, schema); + let spill_file = spill_manager + .spill_record_batch_and_finish(&[sliced_batch], "TestGC")? + .unwrap(); + + // 4. Check file size on disk + let file_size = fs::metadata(spill_file.path().unwrap())?.len(); + + // The original buffer size is around 70KB. + // Without GC, the spill file would be > 70KB. + // With GC, it should be much smaller (only 10 rows of ~70 bytes each + metadata). + assert!( + file_size < 10 * 1024, + "Spill file is too large ({file_size} bytes)! GC might not be working." + ); + + Ok(()) + } + + #[tokio::test] + async fn test_spill_file_size_gc_verification_binary_view() -> Result<()> { + use arrow::array::BinaryViewArray; + use std::fs; + + // 1. Setup bloated data (large buffers) + let num_rows = 1000; + let binary_array: BinaryViewArray = + (0..num_rows).map(|i| Some(vec![i as u8; 100])).collect(); + let schema = Arc::new(Schema::new(vec![Field::new( + "b", + DataType::BinaryView, + false, + )])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(binary_array.clone()) as ArrayRef], + )?; + + // 2. Slice it heavily (1% of the data) + let sliced_batch = batch.slice(0, 10); + + // 3. Spill to disk using SpillManager + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, schema); + let spill_file = spill_manager + .spill_record_batch_and_finish(&[sliced_batch], "TestGCBinary")? + .unwrap(); + + // 4. Check file size on disk + let file_size = fs::metadata(spill_file.path().unwrap())?.len(); + + // Original buffer is 100KB. + // With GC, it should be much smaller. + assert!( + file_size < 10 * 1024, + "Spill file is too large ({file_size} bytes)! GC might not be working." + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/spill/replayable_spill_input.rs b/native/vendor/datafusion-physical-plan/src/spill/replayable_spill_input.rs new file mode 100644 index 0000000..94a0aef --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/spill/replayable_spill_input.rs
@@ -0,0 +1,447 @@ +// 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. + +//! Utility for replaying a one-shot input `RecordBatchStream` through spill. +//! +//! See comments in [`ReplayableStreamSource`] for details. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::SendableRecordBatchStream; +use datafusion_execution::{RecordBatchStream, SpillFile}; +use futures::Stream; +use parking_lot::Mutex; + +use crate::EmptyRecordBatchStream; +use crate::spill::in_progress_spill_file::InProgressSpillFile; +use crate::spill::spill_manager::SpillManager; + +/// Spill-backed replayable stream source. +/// +/// [`ReplayableStreamSource`] is constructed from an input stream, usually produced +/// by executing an input `ExecutionPlan`. +/// +/// - On the first pass, it evaluates the input stream, produces `RecordBatch`es, +/// caches those batches to a local spill file, and also forwards them to the +/// output. +/// - On subsequent passes, it reads directly from the spill file. +/// +/// ```text +/// first pass: +/// +/// RecordBatch stream +/// | +/// v +/// [batch] -> output +/// | +/// +----> spill file +/// +/// +/// later passes: +/// +/// spill file +/// | +/// v +/// [batch] -> output +/// ``` +/// +/// This is useful when an input stream must be replayed and: +/// - Re-evaluation is expensive because the input stream may come from a long +/// and complex pipeline. +/// - The parent operator is under memory pressure and cannot cache the input in +/// memory for replay. +/// +/// # Concurrency assumption +/// Passes must be opened and consumed sequentially. +/// Opening another pass before exhausting the current one returns an error. +pub(crate) struct ReplayableStreamSource { + schema: SchemaRef, + input: Option<SendableRecordBatchStream>, + spill_manager: SpillManager, + request_description: String, + /// Inner state is owned by either the source or one active stream to ensure + /// sequential access; see struct docs for the concurrency contract. + /// + /// Ownership model: + /// - No active stream: source owns the state (`source.state = Some(state)`). + /// - Active stream: the stream owns the state (`source.state = None`). + state: Arc<Mutex<Option<StateInner>>>, +} + +/// Inner state exclusively owned by either [`ReplayableStreamSource`] or one [`ReplayableSpillStream`] +enum StateInner { + Unopened, + Replayable(Option<Arc<dyn SpillFile>>), + Poisoned, +} + +impl ReplayableStreamSource { + /// Creates a replayable stream producer over a one-shot input stream. + /// + /// It caches the input into a local spill file on the first pass, then + /// reads directly from that spill file on subsequent passes. + pub(crate) fn new( + input: SendableRecordBatchStream, + spill_manager: SpillManager, + request_description: impl Into<String>, + ) -> Self { + let schema = input.schema(); + Self { + schema, + input: Some(input), + spill_manager, + request_description: request_description.into(), + state: Arc::new(Mutex::new(Some(StateInner::Unopened))), + } + } + + fn set_state(&self, state: StateInner) { + *self.state.lock() = Some(state); + } + + /// Opens the next pass over this input. + /// + /// The first call returns a stream that forwards upstream batches while + /// caching them to spill. Later calls return streams that read directly + /// from the completed spill file. + /// + /// # Note + /// Subsequent passes MUST be opened only after the previous pass is fully + /// consumed; otherwise, an error is returned. + pub(crate) fn open_pass(&mut self) -> Result<SendableRecordBatchStream> { + let state = self.state.lock().take(); + let Some(state) = state else { + return internal_err!("ReplayableStreamSource pass is still active"); + }; + + match state { + StateInner::Unopened => { + let Some(input) = self.input.take() else { + self.set_state(StateInner::Poisoned); + return internal_err!( + "ReplayableStreamSource missing first-pass input" + ); + }; + let spill_file = match self + .spill_manager + .create_in_progress_file(&self.request_description) + { + Ok(spill_file) => spill_file, + Err(e) => { + self.input = Some(input); + self.set_state(StateInner::Unopened); + return Err(e); + } + }; + + Ok(Box::pin(ReplayableSpillStream::new_first( + Arc::clone(&self.schema), + input, + Arc::clone(&self.state), + spill_file, + ))) + } + StateInner::Poisoned => { + internal_err!( + "ReplayableStreamSource first pass did not complete successfully" + ) + } + StateInner::Replayable(spill_file) => { + let replay_state = spill_file.clone(); + match ReplayableSpillStream::new_replay( + Arc::clone(&self.schema), + &self.spill_manager, + Arc::clone(&self.state), + spill_file, + ) { + Ok(stream) => Ok(Box::pin(stream)), + Err(e) => { + self.set_state(StateInner::Replayable(replay_state)); + Err(e) + } + } + } + } + } +} + +/// Makes a one-shot stream replayable using spill caching, keeping replays fast +/// and memory efficient. +/// +/// On the first pass, it evaluates and forwards output from `inner` while +/// caching it to a spill file for future replays. +/// +/// On later passes, it replays directly from the cached spill file. +/// +/// See also [`ReplayableStreamSource`] for details. +struct ReplayableSpillStream { + schema: SchemaRef, + shared_state: Arc<Mutex<Option<StateInner>>>, + held_state: Option<StateInner>, + spill_file: Option<InProgressSpillFile>, + inner: SendableRecordBatchStream, +} + +impl ReplayableSpillStream { + fn new_first( + schema: SchemaRef, + inner: SendableRecordBatchStream, + shared_state: Arc<Mutex<Option<StateInner>>>, + spill_file: InProgressSpillFile, + ) -> Self { + Self { + schema, + shared_state, + held_state: Some(StateInner::Unopened), + spill_file: Some(spill_file), + inner, + } + } + + fn new_replay( + schema: SchemaRef, + spill_manager: &SpillManager, + shared_state: Arc<Mutex<Option<StateInner>>>, + spill_file: Option<Arc<dyn SpillFile>>, + ) -> Result<Self> { + let inner = if let Some(file) = spill_file.as_ref() { + spill_manager.read_spill_as_stream(Arc::clone(file), None)? + } else { + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) + }; + + Ok(Self { + schema, + shared_state, + held_state: Some(StateInner::Replayable(spill_file)), + spill_file: None, + inner, + }) + } + + fn restore_held_state(&mut self) { + if let Some(state) = self.held_state.take() { + *self.shared_state.lock() = Some(state); + } + } + + fn set_state(&mut self, state: StateInner) { + if self.held_state.take().is_some() { + *self.shared_state.lock() = Some(state); + } + } + + fn poison(&mut self) { + self.set_state(StateInner::Poisoned); + } +} + +impl Stream for ReplayableSpillStream { + type Item = Result<RecordBatch>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let this = self.get_mut(); + + match this.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(batch))) => { + if batch.num_rows() > 0 + && let Some(spill_file) = this.spill_file.as_mut() + && let Err(e) = spill_file.append_batch(&batch) + { + this.spill_file.take(); + this.poison(); + return Poll::Ready(Some(Err(e))); + } + + Poll::Ready(Some(Ok(batch))) + } + Poll::Ready(Some(Err(e))) => { + this.spill_file.take(); + this.poison(); + Poll::Ready(Some(Err(e))) + } + // The stream is exhausted, give the inner state ownership back to `ReplayableStreamSource` + Poll::Ready(None) => { + // Release the input pipeline's resources. + let inner_schema = this.inner.schema(); + this.inner = Box::pin(EmptyRecordBatchStream::new(inner_schema)); + if let Some(spill_file) = this.spill_file.as_mut() { + match spill_file.finish() { + Ok(file) => { + this.spill_file.take(); + this.set_state(StateInner::Replayable(file)); + Poll::Ready(None) + } + Err(e) => { + this.spill_file.take(); + this.poison(); + Poll::Ready(Some(Err(e))) + } + } + } else { + this.restore_held_state(); + Poll::Ready(None) + } + } + Poll::Pending => Poll::Pending, + } + } +} + +impl RecordBatchStream for ReplayableSpillStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Drop for ReplayableSpillStream { + /// If a stream is dropped before it finishes, poison the state so later + /// replay attempts fail. + /// + /// A partial first pass leaves the spill file incomplete, so replaying it + /// would be unsafe. + fn drop(&mut self) { + if self.held_state.is_some() { + self.poison(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + use futures::{StreamExt, TryStreamExt}; + + use crate::stream::RecordBatchStreamAdapter; + + fn build_spill_manager(schema: SchemaRef) -> Result<SpillManager> { + let runtime = Arc::new(RuntimeEnvBuilder::new().build()?); + let metrics_set = ExecutionPlanMetricsSet::new(); + let spill_metrics = SpillMetrics::new(&metrics_set, 0); + Ok(SpillManager::new(runtime, spill_metrics, schema)) + } + + fn build_batch(schema: SchemaRef, values: Vec<i64>) -> Result<RecordBatch> { + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(values))]) + .map_err(Into::into) + } + + #[tokio::test] + async fn test_replayable_spill_input_replays_completed_first_pass() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let batch1 = build_batch(Arc::clone(&schema), vec![1, 2])?; + let batch2 = build_batch(Arc::clone(&schema), vec![3, 4])?; + + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch1.clone()), Ok(batch2.clone())]), + )); + let spill_manager = build_spill_manager(Arc::clone(&schema))?; + let mut replayable = + ReplayableStreamSource::new(input, spill_manager, "test replayable spill"); + + let pass1 = replayable.open_pass()?; + let pass1_batches = pass1.try_collect::<Vec<_>>().await?; + assert_eq!(pass1_batches, vec![batch1.clone(), batch2.clone()]); + + let pass2 = replayable.open_pass()?; + let pass2_batches = pass2.try_collect::<Vec<_>>().await?; + assert_eq!(pass2_batches, vec![batch1, batch2]); + + Ok(()) + } + + // Try to open a new pass, when the first pass has not finished. + // The spill file is only partially written, so an error will be returned. + #[tokio::test] + async fn test_replayable_spill_input_poisoned_when_first_pass_dropped() -> Result<()> + { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let batch1 = build_batch(Arc::clone(&schema), vec![1, 2])?; + let batch2 = build_batch(Arc::clone(&schema), vec![3, 4])?; + + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch1), Ok(batch2)]), + )); + let spill_manager = build_spill_manager(Arc::clone(&schema))?; + let mut replayable = + ReplayableStreamSource::new(input, spill_manager, "test replayable spill"); + + let mut pass1 = replayable.open_pass()?; + let first = pass1.next().await.transpose()?; + assert!(first.is_some()); + drop(pass1); + + let err = match replayable.open_pass() { + Ok(_) => panic!("expected first pass to poison replayable spill input"), + Err(err) => err.strip_backtrace(), + }; + assert!( + err.to_string().contains( + "ReplayableStreamSource first pass did not complete successfully" + ) + ); + + Ok(()) + } + + // Open a new pass, when the previous pass from spill is still in progress. + // An error is expected, since it requires sequential access. + #[tokio::test] + async fn test_replayable_spill_input_errors_when_replay_pass_in_progress() + -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let batch1 = build_batch(Arc::clone(&schema), vec![1, 2])?; + let batch2 = build_batch(Arc::clone(&schema), vec![3, 4])?; + + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch1.clone()), Ok(batch2.clone())]), + )); + let spill_manager = build_spill_manager(Arc::clone(&schema))?; + let mut replayable = + ReplayableStreamSource::new(input, spill_manager, "test replayable spill"); + + let pass1 = replayable.open_pass()?; + let _ = pass1.try_collect::<Vec<_>>().await?; + + let pass2 = replayable.open_pass()?; + let err = match replayable.open_pass() { + Ok(_) => panic!("expected open_pass to fail while replay pass is active"), + Err(err) => err.strip_backtrace(), + }; + assert!( + err.to_string() + .contains("ReplayableStreamSource pass is still active") + ); + drop(pass2); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/spill/spill_manager.rs b/native/vendor/datafusion-physical-plan/src/spill/spill_manager.rs new file mode 100644 index 0000000..aee9e91 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/spill/spill_manager.rs
@@ -0,0 +1,409 @@ +// 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. + +//! Define the `SpillManager` struct, which is responsible for reading and writing `RecordBatch`es to raw files based on the provided configurations. + +use super::{SpillReaderStream, in_progress_spill_file::InProgressSpillFile}; +use crate::coop::cooperative; +use crate::{common::spawn_buffered, metrics::SpillMetrics}; +use arrow::array::{BinaryViewArray, GenericByteViewArray, StringViewArray}; +use arrow::datatypes::{ByteViewType, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result, config::SpillCompression}; +use datafusion_execution::SendableRecordBatchStream; +use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_execution::spill_file::SpillFile; +use std::borrow::Borrow; +use std::sync::Arc; + +/// The `SpillManager` is responsible for the following tasks: +/// - Reading and writing `RecordBatch`es to raw files based on the provided configurations. +/// - Updating the associated metrics. +/// +/// Note: The caller (external operators such as `SortExec`) is responsible for interpreting the spilled files. +/// For example, all records within the same spill file are ordered according to a specific order. +#[derive(Debug, Clone)] +pub struct SpillManager { + env: Arc<RuntimeEnv>, + pub(crate) metrics: SpillMetrics, + schema: SchemaRef, + /// Number of batches to buffer in memory during disk reads + batch_read_buffer_capacity: usize, + /// general-purpose compression options + pub(crate) compression: SpillCompression, +} + +impl SpillManager { + pub fn new(env: Arc<RuntimeEnv>, metrics: SpillMetrics, schema: SchemaRef) -> Self { + Self { + env, + metrics, + schema, + batch_read_buffer_capacity: 2, + compression: SpillCompression::default(), + } + } + + pub fn with_batch_read_buffer_capacity( + mut self, + batch_read_buffer_capacity: usize, + ) -> Self { + self.batch_read_buffer_capacity = batch_read_buffer_capacity; + self + } + + pub fn with_compression_type(mut self, spill_compression: SpillCompression) -> Self { + self.compression = spill_compression; + self + } + + /// Returns the schema for batches managed by this SpillManager + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + pub(crate) fn env(&self) -> &RuntimeEnv { + &self.env + } + + /// Creates a temporary file for in-progress operations, returning an error + /// message if file creation fails. The file can be used to append batches + /// incrementally and then finish the file when done. + pub fn create_in_progress_file( + &self, + request_msg: &str, + ) -> Result<InProgressSpillFile> { + let temp_file = self.env.disk_manager.create_tmp_file(request_msg)?; + Ok(InProgressSpillFile::new(Arc::new(self.clone()), temp_file)) + } + + /// Spill input `batches` into a single file in a atomic operation. If it is + /// intended to incrementally write in-memory batches into the same spill file, + /// use [`Self::create_in_progress_file`] instead. + /// None is returned if no batches are spilled. + /// + /// # Errors + /// - Returns an error if spilling would exceed the disk usage limit configured + /// by `max_temp_directory_size` in `DiskManager` + pub fn spill_record_batch_and_finish( + &self, + batches: &[RecordBatch], + request_msg: &str, + ) -> Result<Option<Arc<dyn SpillFile>>> { + let mut in_progress_file = self.create_in_progress_file(request_msg)?; + + for batch in batches { + in_progress_file.append_batch(batch)?; + } + + in_progress_file.finish() + } + + /// Spill an iterator of `RecordBatch`es to disk and return the spill file and the size of the largest batch in memory + /// Note that this expects the caller to provide *non-sliced* batches, so the memory calculation of each batch is accurate. + pub(crate) fn spill_record_batch_iter_and_return_max_batch_memory( + &self, + mut iter: impl Iterator<Item = Result<impl Borrow<RecordBatch>>>, + request_description: &str, + ) -> Result<Option<(Arc<dyn SpillFile>, usize)>> { + let mut in_progress_file = self.create_in_progress_file(request_description)?; + + let mut max_record_batch_size = 0; + + iter.try_for_each(|batch| { + let batch = batch?; + let borrowed = batch.borrow(); + if borrowed.num_rows() == 0 { + return Ok(()); + } + let gc_sliced_size = in_progress_file.append_batch(borrowed)?; + max_record_batch_size = max_record_batch_size.max(gc_sliced_size); + Result::<_, DataFusionError>::Ok(()) + })?; + + let file = in_progress_file.finish()?; + + Ok(file.map(|f| (f, max_record_batch_size))) + } + + /// Spill a stream of `RecordBatch`es to disk and return the spill file and the size of the largest batch in memory + pub(crate) async fn spill_record_batch_stream_and_return_max_batch_memory( + &self, + stream: &mut SendableRecordBatchStream, + request_description: &str, + ) -> Result<Option<(Arc<dyn SpillFile>, usize)>> { + use futures::StreamExt; + + let mut in_progress_file = self.create_in_progress_file(request_description)?; + + let mut max_record_batch_size = 0; + + while let Some(batch) = stream.next().await { + let batch = batch?; + let gc_sliced_size = in_progress_file.append_batch(&batch)?; + + max_record_batch_size = max_record_batch_size.max(gc_sliced_size); + } + + let file = in_progress_file.finish()?; + + Ok(file.map(|f| (f, max_record_batch_size))) + } + + /// Reads a spill file as a stream. The file must be created by the current + /// `SpillManager`; otherwise an error will be returned. + /// + /// Output is produced in FIFO order: the batch appended first is read first. + /// + /// # Arg `max_record_batch_memory` + /// + /// Most callers should pass `None`. This is mainly useful for the + /// memory-limited sort-preserving merge path. + /// + /// When provided, this value is used only as a validation hint. If a + /// decoded batch exceeds this threshold, a debug-level log message is + /// emitted. + /// + /// That path uses the maximum spilled batch size to conservatively estimate + /// the merge degree when merging multiple sorted runs. + pub fn read_spill_as_stream( + &self, + spill_file_path: Arc<dyn SpillFile>, + max_record_batch_memory: Option<usize>, + ) -> Result<SendableRecordBatchStream> { + let stream = Box::pin(cooperative(SpillReaderStream::new( + Arc::clone(&self.schema), + spill_file_path, + max_record_batch_memory, + )?)); + + Ok(spawn_buffered(stream, self.batch_read_buffer_capacity)) + } + + /// Same as `read_spill_as_stream`, but without buffering. + pub fn read_spill_as_stream_unbuffered( + &self, + spill_file_path: Arc<dyn SpillFile>, + max_record_batch_memory: Option<usize>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(cooperative(SpillReaderStream::new( + Arc::clone(&self.schema), + spill_file_path, + max_record_batch_memory, + )?))) + } +} + +pub(crate) trait GetSlicedSize { + /// Returns the size of the `RecordBatch` when sliced. + /// Note: if multiple arrays or even a single array share the same data buffers, we may double count each buffer. + /// Therefore, make sure we call gc() or gc_view_arrays() before using this method. + fn get_sliced_size(&self) -> Result<usize>; +} + +impl GetSlicedSize for RecordBatch { + fn get_sliced_size(&self) -> Result<usize> { + let mut total = 0; + for array in self.columns() { + let data = array.to_data(); + total += data.get_slice_memory_size()?; + + // While StringViewArray holds large data buffer for non inlined string, the Arrow layout (BufferSpec) + // does not include any data buffers. Currently, ArrayData::get_slice_memory_size() + // under-counts memory size by accounting only views buffer although data buffer is cloned during slice() + // + // Therefore, we manually add the sum of the lengths used by all non inlined views + // on top of the sliced size for views buffer. This matches the intended semantics of + // "bytes needed if we materialized exactly this slice into fresh buffers". + // This is a workaround until https://github.com/apache/arrow-rs/issues/8230 + if let Some(sv) = array.as_any().downcast_ref::<StringViewArray>() { + total += byte_view_data_buffer_size(sv); + } + if let Some(bv) = array.as_any().downcast_ref::<BinaryViewArray>() { + total += byte_view_data_buffer_size(bv); + } + } + Ok(total) + } +} + +fn byte_view_data_buffer_size<T: ByteViewType>(array: &GenericByteViewArray<T>) -> usize { + array + .data_buffers() + .iter() + .map(|buffer| buffer.capacity()) + .sum() +} + +#[cfg(test)] +mod tests { + use super::SpillManager; + use crate::common::collect; + use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; + use crate::spill::{get_record_batch_memory_size, spill_manager::GetSlicedSize}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::{ + array::{ArrayRef, Int32Array, StringArray, StringViewArray}, + record_batch::RecordBatch, + }; + use datafusion_common::Result; + use datafusion_execution::runtime_env::RuntimeEnv; + use std::sync::Arc; + + fn build_test_spill_manager( + env: Arc<RuntimeEnv>, + schema: Arc<Schema>, + ) -> SpillManager { + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + SpillManager::new(env, metrics, schema) + } + + fn build_writer_batch(schema: Arc<Schema>) -> Result<RecordBatch> { + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .map_err(Into::into) + } + + #[tokio::test] + async fn test_read_spill_as_stream_from_another_spill_manager_same_schema() + -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let writer_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let reader_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + + let writer = + build_test_spill_manager(Arc::clone(&env), Arc::clone(&writer_schema)); + let reader = build_test_spill_manager(env, Arc::clone(&reader_schema)); + let written_batch = build_writer_batch(Arc::clone(&writer_schema))?; + + let spill_file = writer + .spill_record_batch_and_finish( + std::slice::from_ref(&written_batch), + "writer", + )? + .unwrap(); + + // Same-schema reads through a different SpillManager currently pass + // because only schema compatibility is validated. This is not a + // supported usage pattern. + let stream = reader.read_spill_as_stream(spill_file, None)?; + assert_eq!(stream.schema(), reader_schema); + + let batches = collect(stream).await?; + assert_eq!(batches, vec![written_batch]); + + Ok(()) + } + + #[tokio::test] + async fn test_read_spill_as_stream_from_another_spill_manager_different_schema() + -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let writer_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let reader_schema = Arc::new(Schema::new(vec![ + Field::new("other_id", DataType::Int32, true), + Field::new("other_value", DataType::Utf8, true), + ])); + + let writer = + build_test_spill_manager(Arc::clone(&env), Arc::clone(&writer_schema)); + let reader = build_test_spill_manager(env, Arc::clone(&reader_schema)); + let written_batch = build_writer_batch(Arc::clone(&writer_schema))?; + + let spill_file = writer + .spill_record_batch_and_finish( + std::slice::from_ref(&written_batch), + "writer", + )? + .unwrap(); + + let stream = reader.read_spill_as_stream(spill_file, None)?; + let err = collect(stream) + .await + .expect_err("schema mismatch should fail fast"); + let err = err.to_string(); + assert!(err.contains("Spill file schema mismatch")); + assert!(err.contains("expected")); + assert!(err.contains("got")); + + Ok(()) + } + + #[test] + fn check_sliced_size_for_string_view_array() -> Result<()> { + let array_length = 50; + let short_len = 8; + let long_len = 25; + + // Build StringViewArray that includes both inline strings and non inlined strings + let strings: Vec<String> = (0..array_length) + .map(|i| { + if i % 2 == 0 { + "a".repeat(short_len) + } else { + "b".repeat(long_len) + } + }) + .collect(); + + let string_array = StringViewArray::from(strings); + let array_ref: ArrayRef = Arc::new(string_array); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "strings", + DataType::Utf8View, + false, + )])), + vec![array_ref], + ) + .unwrap(); + + // We did not slice the batch, so these two memory size should be equal + assert_eq!( + batch.get_sliced_size().unwrap(), + get_record_batch_memory_size(&batch) + ); + + // Slice the batch into half + let half_batch = batch.slice(0, array_length / 2); + // Now sliced_size is smaller because the views buffer is sliced + assert!( + half_batch.get_sliced_size().unwrap() + < get_record_batch_memory_size(&half_batch) + ); + let data = arrow::array::Array::to_data(&half_batch.column(0)); + let views_sliced_size = data.get_slice_memory_size()?; + // The sliced size should be larger than sliced views buffer size + assert!(views_sliced_size < half_batch.get_sliced_size().unwrap()); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/spill/spill_pool.rs b/native/vendor/datafusion-physical-plan/src/spill/spill_pool.rs new file mode 100644 index 0000000..6e964d7 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/spill/spill_pool.rs
@@ -0,0 +1,1648 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use futures::{Stream, StreamExt}; +use std::collections::VecDeque; +use std::mem; +use std::sync::Arc; +use std::task::Waker; + +use parking_lot::Mutex; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, SpillFile}; + +use super::in_progress_spill_file::InProgressSpillFile; +use super::spill_manager::SpillManager; + +/// Shared state between the writer and readers of a spill pool. +/// This contains the queue of files and coordination state. +/// +/// # Locking Design +/// +/// This struct uses **fine-grained locking** with nested `Arc<Mutex<>>`: +/// - `SpillPoolShared` is wrapped in `Arc<Mutex<>>` (outer lock) +/// - Each `ActiveSpillFileShared` is wrapped in `Arc<Mutex<>>` (inner lock) +/// +/// This enables: +/// 1. **Short critical sections**: The outer lock is held only for queue operations +/// 2. **I/O outside locks**: Disk I/O happens while holding only the file-specific lock +/// 3. **Concurrent operations**: Reader can access the queue while writer does I/O +/// +/// **Lock ordering discipline**: Never hold both locks simultaneously to prevent deadlock. +/// Always: acquire outer lock → release outer lock → acquire inner lock (if needed). +struct SpillPoolShared { + /// Queue of ALL files (including the current write files if any exist). + /// Readers always read from the front of this queue (FIFO). + /// Each file has its own lock to enable concurrent reader/writer access. + files: VecDeque<Arc<Mutex<ActiveSpillFileShared>>>, + /// SpillManager for creating files and tracking metrics + spill_manager: Arc<SpillManager>, + /// Pool-level waker to notify when new files are available (single reader) + waker: Option<Waker>, + /// FIFO queue of open write files. The queue may contain multiple items when multiple + /// writers concurrently write to the pool. + /// Each write file has its own lock to allow I/O without blocking queue access. + open_write_files: VecDeque<Arc<Mutex<ActiveSpillFileShared>>>, + /// Number of `SpillPoolWriter` instances that have not been dropped yet. As long as this value + /// is greater than zero, readers should assume batches may still be pushed. This prevents + /// premature EOF signaling. + remaining_writer_count: usize, +} + +impl SpillPoolShared { + /// Creates a new shared pool state + fn new(spill_manager: Arc<SpillManager>) -> Self { + Self { + files: VecDeque::new(), + spill_manager, + waker: None, + open_write_files: VecDeque::new(), + remaining_writer_count: 1, + } + } + + /// Registers a waker to be notified when new data is available (pool-level) + fn register_waker(&mut self, waker: Waker) { + self.waker = Some(waker); + } + + /// Wakes the pool-level reader + fn wake(&mut self) { + if let Some(waker) = self.waker.take() { + waker.wake(); + } + } +} + +/// Writer for a spill pool that can be cloned to produce additional writers. +/// +/// Created by [`mpsc_channel`]. See that function for architecture diagrams and usage +/// examples. +pub struct SpillPoolWriter { + /// The underlying shared writer. Kept private and never cloned, so this pool always has + /// exactly one writer. + inner: SpillPoolSink, +} + +impl SpillPoolWriter { + /// Spills a batch to the pool, rotating files when necessary. + /// + /// See [`mpsc_channel`] for the rotation semantics. + /// + /// # Errors + /// + /// Returns an error if disk I/O fails or disk quota is exceeded. + pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { + self.inner.push_batch(batch) + } +} + +impl SpillPoolWriter { + /// Returns a new sink that can be used to spill batches to the pool. + /// + /// As an alternative to this function, it is also possible to clone the writer. The benefit + /// of this method is that the output type matches the type used by [`spsc_channel`]. This + /// enables cost-free abstraction for producers over SPSC and MPSC channels. + pub fn new_sink(&self) -> SpillPoolSink { + // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` + // implementation of `SpillPoolWriter`. + self.inner.shared.lock().remaining_writer_count += 1; + SpillPoolSink { + max_file_size_bytes: self.inner.max_file_size_bytes, + shared: Arc::clone(&self.inner.shared), + } + } +} + +impl Clone for SpillPoolWriter { + fn clone(&self) -> Self { + Self { + inner: self.new_sink(), + } + } +} + +impl Drop for SpillPoolSink { + fn drop(&mut self) { + let mut shared = self.shared.lock(); + + shared.remaining_writer_count -= 1; + let is_last_writer = shared.remaining_writer_count == 0; + + if !is_last_writer { + // Other writer clones are still active; do not finalize or + // signal EOF to readers. + return; + } + + // Finalize any spill files that were not finished yet + if !shared.open_write_files.is_empty() { + let files = mem::take(&mut shared.open_write_files); + drop(shared); + + for file in files { + let mut file_shared = file.lock(); + + // Finish the current writer if it exists + if let Some(mut writer) = file_shared.writer.take() { + // Ignore errors on drop - we're in destructor + let _ = writer.finish(); + } + + // Mark as finished so readers know not to wait for more data + file_shared.writer_finished = true; + + // Wake reader waiting on this file (it's now finished) + file_shared.wake(); + drop(file_shared); + } + + shared = self.shared.lock(); + } + + // Wake pool-level readers + shared.wake(); + } +} + +/// Single writer for a spill pool that cannot be cloned. +/// +/// Created by [`spsc_channel`] and [`SpillPoolWriter::new_sink`]. +pub struct SpillPoolSink { + /// Maximum size in bytes before rotating to a new file. + /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`. + max_file_size_bytes: usize, + /// Shared state with readers (includes current_write_file for coordination) + shared: Arc<Mutex<SpillPoolShared>>, +} + +impl SpillPoolSink { + /// Spills a batch to the pool, rotating files when necessary. + /// + /// See [`spsc_channel`] for overall architecture and examples. + /// + /// # Errors + /// + /// Returns an error if disk I/O fails or disk quota is exceeded. + pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { + if batch.num_rows() == 0 { + // Skip empty batches + return Ok(()); + } + + let batch_size = batch.get_array_memory_size(); + + // Fine-grained locking: Lock shared state briefly for queue access + let mut shared = self.shared.lock(); + + // Create new file if there is none available to append to + let write_file = if !shared.open_write_files.is_empty() { + shared.open_write_files.pop_front().unwrap() + } else { + let spill_manager = Arc::clone(&shared.spill_manager); + // Release shared lock before disk I/O (fine-grained locking) + drop(shared); + + let writer = spill_manager.create_in_progress_file("SpillPool")?; + // Clone the file so readers can access it immediately + let file = Arc::clone(writer.file().expect( + "InProgressSpillFile should always have a file when it is first created", + )); + + let file_shared = Arc::new(Mutex::new(ActiveSpillFileShared { + writer: Some(writer), + file: Some(file), // Set immediately so readers can access it + batches_written: 0, + estimated_size: 0, + writer_finished: false, + waker: None, + })); + + // Re-acquire lock and push to shared queue + shared = self.shared.lock(); + shared.files.push_back(Arc::clone(&file_shared)); + shared.wake(); // Wake readers waiting for new files + file_shared + }; + + // Release shared lock before file I/O (fine-grained locking) + // This allows readers to access the queue while we do disk I/O + drop(shared); + + // Write batch to current file - lock only the specific file + let mut file_shared = write_file.lock(); + + // Append the batch + if let Some(ref mut writer) = file_shared.writer { + writer.append_batch(batch)?; + // make sure we flush the writer for readers + writer.flush()?; + file_shared.batches_written += 1; + file_shared.estimated_size += batch_size; + } + + // Wake reader waiting on this specific file + file_shared.wake(); + + let max_file_size_reached = file_shared.estimated_size > self.max_file_size_bytes; + + if max_file_size_reached { + // Finish the IPC writer + if let Some(mut writer) = file_shared.writer.take() { + writer.finish()?; + } + // Mark as finished so readers know not to wait for more data + file_shared.writer_finished = true; + // Wake reader waiting on this file (it's now finished) + file_shared.wake(); + + // Don't place `write_file` back in the `open_write_files` queue so we don't + // try writing to it again + } else { + // Release file lock + drop(file_shared); + // Put back the current file for further writing + let mut shared = self.shared.lock(); + shared.open_write_files.push_back(write_file); + } + + Ok(()) + } +} + +/// Creates a paired writer and reader for a spill pool with SPSC (single-producer, +/// single-consumer) semantics and strict FIFO ordering. +/// +/// If you need a spill pool that supports several producers, use [`mpsc_channel`] instead. +/// +/// The reader can start reading immediately after the writer appends a batch +/// to the spill file, without waiting for the file to be sealed, while the writer continues to +/// write more data. +/// +/// Internally this coordinates rotating spill files based on size limits, and +/// handles asynchronous notification between the writer and reader using wakers. +/// This ensures that we manage disk usage efficiently while allowing concurrent +/// I/O between the writer and reader. +/// +/// # Data Flow Overview +/// +/// 1. Writer write batch `B0` to F1 +/// 2. Writer write batch `B1` to F1, notices the size limit exceeded, finishes F1. +/// 3. Reader read `B0` from F1 +/// 4. Reader read `B1`, no more batch to read -> wait on the waker +/// 5. Writer write batch `B2` to a new file `F2`, wake up the waiting reader. +/// 6. Reader read `B2` from F2. +/// 7. Repeat until writer is dropped. +/// +/// # Architecture +/// +/// ```text +/// ┌─────────────────────────────────────────────────────────────────────────┐ +/// │ SpillPool │ +/// │ │ +/// │ Writer Side Shared State Reader Side │ +/// │ ─────────── ──────────── ─────────── │ +/// │ │ +/// │ SpillPoolSink ┌────────────────────┐ RecordBatchStream │ +/// │ │ │ VecDeque<File> │ │ │ +/// │ │ │ ┌────┐┌────┐ │ │ │ +/// │ push_batch() │ │ F1 ││ F2 │ ... │ next().await │ +/// │ │ │ └────┘└────┘ │ │ │ +/// │ ▼ │ │ ▼ │ +/// │ ┌─────────┐ │ │ ┌──────────┐ │ +/// │ │Current │───────▶│ Coordination: │◀───│ Current │ │ +/// │ │Write │ │ - Wakers │ │ Read │ │ +/// │ │File │ │ - Batch counts │ │ File │ │ +/// │ └─────────┘ │ - Writer status │ └──────────┘ │ +/// │ │ └────────────────────┘ │ │ +/// │ │ │ │ +/// │ Size > limit? Read all batches? │ +/// │ │ │ │ +/// │ ▼ ▼ │ +/// │ Rotate to new file Pop from queue │ +/// └─────────────────────────────────────────────────────────────────────────┘ +/// +/// Writer produces → Shared queue → Reader consumes +/// ``` +/// +/// # File State Machine +/// +/// Each file in the pool coordinates between writer and reader: +/// +/// ```text +/// Writer View Reader View +/// ─────────── ─────────── +/// +/// Created writer: Some(..) batches_read: 0 +/// batches_written: 0 (waiting for data) +/// │ +/// ▼ +/// Writing append_batch() Can read if: +/// batches_written++ batches_read < batches_written +/// wake readers +/// │ │ +/// │ ▼ +/// ┌──────┴──────┐ poll_next() → batch +/// │ │ batches_read++ +/// ▼ ▼ +/// Size > limit? More data? +/// │ │ +/// │ └─▶ Yes ──▶ Continue writing +/// ▼ +/// finish() Reader catches up: +/// writer_finished = true batches_read == batches_written +/// wake readers │ +/// │ ▼ +/// └─────────────────────▶ Returns Poll::Ready(None) +/// File complete, pop from queue +/// ``` +/// +/// # Arguments +/// +/// * `max_file_size_bytes` - Maximum size per file before rotation. When a file +/// exceeds this size, the writer automatically rotates to a new file. +/// * `spill_manager` - Manager for file creation and metrics tracking +/// +/// # Returns +/// +/// A tuple of `(SpillPoolSink, SendableRecordBatchStream)` that share the same +/// underlying pool. The reader is returned as a stream for immediate use with +/// async stream combinators. +/// +/// # Example +/// +/// ``` +/// use std::sync::Arc; +/// use arrow::array::{ArrayRef, Int32Array}; +/// use arrow::datatypes::{DataType, Field, Schema}; +/// use arrow::record_batch::RecordBatch; +/// use datafusion_execution::runtime_env::RuntimeEnv; +/// use futures::StreamExt; +/// +/// # use datafusion_physical_plan::spill::spill_pool; +/// # use datafusion_physical_plan::spill::SpillManager; // Re-exported for doctests +/// # use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; +/// # +/// # #[tokio::main] +/// # async fn main() -> datafusion_common::Result<()> { +/// # // Setup for the example (typically comes from TaskContext in production) +/// # let env = Arc::new(RuntimeEnv::default()); +/// # let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); +/// # let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); +/// # let spill_manager = Arc::new(SpillManager::new(env, metrics, schema.clone())); +/// # +/// // Create channel with 1MB file size limit +/// let (writer, mut reader) = spill_pool::spsc_channel(1024 * 1024, spill_manager); +/// +/// // Spawn writer and reader concurrently; writer wakes reader via wakers +/// let writer_task = tokio::spawn(async move { +/// for i in 0..5 { +/// let array: ArrayRef = Arc::new(Int32Array::from(vec![i; 100])); +/// let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); +/// writer.push_batch(&batch)?; +/// } +/// // Explicitly drop writer to finalize the spill file and wake the reader +/// drop(writer); +/// datafusion_common::Result::<()>::Ok(()) +/// }); +/// +/// let reader_task = tokio::spawn(async move { +/// let mut batches_read = 0; +/// while let Some(result) = reader.next().await { +/// let _batch = result?; +/// batches_read += 1; +/// } +/// datafusion_common::Result::<usize>::Ok(batches_read) +/// }); +/// +/// let (writer_res, reader_res) = tokio::join!(writer_task, reader_task); +/// writer_res +/// .map_err(|e| datafusion_common::DataFusionError::Execution(e.to_string()))??; +/// let batches_read = reader_res +/// .map_err(|e| datafusion_common::DataFusionError::Execution(e.to_string()))??; +/// +/// assert_eq!(batches_read, 5); +/// # Ok(()) +/// # } +/// ``` +/// +/// # Why rotate files? +/// +/// File rotation ensures we don't end up with unreferenced disk usage. +/// If we used a single file for all spilled data, we would end up with +/// unreferenced data at the beginning of the file that has already been read +/// by readers but we can't delete because you can't truncate from the start of a file. +/// +/// Consider the case of a query like `SELECT * FROM large_table WHERE false`. +/// Obviously this query produces no output rows, but if we had a spilling operator +/// in the middle of this query between the scan and the filter it would see the entire +/// `large_table` flow through it and thus would spill all of that data to disk. +/// So we'd end up using up to `size(large_table)` bytes of disk space. +/// If instead we use file rotation, and as long as the readers can keep up with the writer, +/// then we can ensure that once a file is fully read by all readers it can be deleted, +/// thus bounding the maximum disk usage to roughly `max_file_size_bytes`. +pub fn spsc_channel( + max_file_size_bytes: usize, + spill_manager: Arc<SpillManager>, +) -> (SpillPoolSink, SendableRecordBatchStream) { + let schema = Arc::clone(spill_manager.schema()); + let shared = Arc::new(Mutex::new(SpillPoolShared::new(spill_manager))); + + let writer = SpillPoolSink { + max_file_size_bytes, + shared: Arc::clone(&shared), + }; + + let reader = SpillPoolReader::new(shared, schema); + + (writer, Box::pin(reader)) +} + +/// Alias for [`mpsc_channel`]. +#[deprecated(note = "Use mpsc_channel instead")] +pub fn channel( + max_file_size_bytes: usize, + spill_manager: Arc<SpillManager>, +) -> (SpillPoolWriter, SendableRecordBatchStream) { + mpsc_channel(max_file_size_bytes, spill_manager) +} + +/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, +/// single-consumer) semantics. See [`spsc_channel`] for the general architecture description +/// of the spill pool. +/// +/// Additional writers can be created by cloning the returned [`SpillPoolWriter`]. +/// +/// In contrast to [`spsc_channel`], this implementation provides no guarantees regarding +/// the read order of the returned [`SendableRecordBatchStream`]. +/// +/// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact +/// write order), use [`spsc_channel`] instead. +/// +/// # File Management +/// +/// The shared channel uses the same size-based rotation trigger as the [single producer channel](spsc_channel). +/// All writers share the same pool of write files and coordinate file rotation. The number of open +/// files is kept as small as possible. When more writes occur concurrently than there are open write +/// files an additional file will be opened to write to. This prevents multiple writers from blocking +/// each other. +/// +/// When the last writer clone is dropped, it finalizes any remaining open write files so that all +/// written data can be accessed by the reader. +/// +/// # Returns +/// +/// A tuple of `(SpillPoolWriter, SendableRecordBatchStream)` that share the same +/// underlying pool. The reader is returned as a stream for immediate use with +/// async stream combinators. The writer can be cloned to create additional writers. +pub fn mpsc_channel( + max_file_size_bytes: usize, + spill_manager: Arc<SpillManager>, +) -> (SpillPoolWriter, SendableRecordBatchStream) { + let (inner, reader) = spsc_channel(max_file_size_bytes, spill_manager); + (SpillPoolWriter { inner }, reader) +} + +/// Shared state between writer and readers for an active spill file. +/// Protected by a Mutex to coordinate between concurrent readers and the writer. +struct ActiveSpillFileShared { + /// Writer handle - taken (set to None) when finish() is called + writer: Option<InProgressSpillFile>, + /// The spill file, set when the writer finishes. + /// Taken by the reader when creating a stream (the file stays open via file handles). + file: Option<Arc<dyn SpillFile>>, + /// Total number of batches written to this file + batches_written: usize, + /// Estimated size in bytes of data written to this file + estimated_size: usize, + /// Whether the writer has finished writing to this file + writer_finished: bool, + /// Waker for reader waiting on this specific file (SPSC: only one reader) + waker: Option<Waker>, +} + +impl ActiveSpillFileShared { + /// Registers a waker to be notified when new data is written to this file + fn register_waker(&mut self, waker: Waker) { + self.waker = Some(waker); + } + + /// Wakes the reader waiting on this file + fn wake(&mut self) { + if let Some(waker) = self.waker.take() { + waker.wake(); + } + } +} + +/// Reader state for a SpillPoolFile (owned by individual SpillPoolFile instances). +/// This is kept separate from the shared state to avoid holding locks during I/O. +struct SpillPoolFileReader { + /// The actual stream reading from disk + stream: SendableRecordBatchStream, + /// Number of batches this reader has consumed + batches_read: usize, +} + +struct SpillPoolFile { + /// Shared coordination state (contains writer and batch counts) + shared: Arc<Mutex<ActiveSpillFileShared>>, + /// Reader state (lazy-initialized, owned by this SpillPoolFile) + reader: Option<SpillPoolFileReader>, + /// Spill manager for creating readers + spill_manager: Arc<SpillManager>, +} + +impl Stream for SpillPoolFile { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll<Option<Self::Item>> { + use std::task::Poll; + + // Step 1: Lock shared state and check coordination + let (should_read, file) = { + let mut shared = self.shared.lock(); + + // Determine if we can read + let batches_read = self.reader.as_ref().map_or(0, |r| r.batches_read); + + if batches_read < shared.batches_written { + // More data available to read - take the file if we don't have a reader yet + let file = if self.reader.is_none() { + shared.file.take() + } else { + None + }; + (true, file) + } else if shared.writer_finished { + // No more data and writer is done - EOF + return Poll::Ready(None); + } else { + // Caught up to writer, but writer still active - register waker and wait + shared.register_waker(cx.waker().clone()); + return Poll::Pending; + } + }; // Lock released here + + // Step 2: Lazy-create reader stream if needed + if self.reader.is_none() && should_read { + if let Some(file) = file { + // we want this unbuffered because files are actively being written to + match self + .spill_manager + .read_spill_as_stream_unbuffered(file, None) + { + Ok(stream) => { + self.reader = Some(SpillPoolFileReader { + stream, + batches_read: 0, + }); + } + Err(e) => return Poll::Ready(Some(Err(e))), + } + } else { + // File not available yet (writer hasn't finished or already taken) + // Register waker and wait for file to be ready + let mut shared = self.shared.lock(); + shared.register_waker(cx.waker().clone()); + return Poll::Pending; + } + } + + // Step 3: Poll the reader stream (no lock held) + if let Some(reader) = &mut self.reader { + match reader.stream.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + // Successfully read a batch - increment counter + reader.batches_read += 1; + Poll::Ready(Some(Ok(batch))) + } + Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), + Poll::Ready(None) => { + // Stream exhausted unexpectedly + // This shouldn't happen if coordination is correct, but handle gracefully + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } else { + // Should not reach here, but handle gracefully + Poll::Ready(None) + } + } +} + +/// A stream that reads from a SpillPool. The reader guarantees FIFO order if a single writer is used. +/// +/// Created by [`spsc_channel`]. See that function for architecture diagrams and usage examples. +/// +/// The stream automatically handles file rotation and reads from completed files. +/// When no data is available, it returns `Poll::Pending` and registers a waker to +/// be notified when the writer produces more data. +/// +/// # Infinite Stream Semantics +/// +/// This stream never returns `None` (`Poll::Ready(None)`) on its own - it will keep +/// waiting for the writer to produce more data. The stream ends only when: +/// - The reader is dropped +/// - The writer is dropped AND all queued data has been consumed +/// +/// This makes it suitable for continuous streaming scenarios where the writer may +/// produce data intermittently. +pub struct SpillPoolReader { + /// Shared reference to the spill pool + shared: Arc<Mutex<SpillPoolShared>>, + /// Current SpillPoolFile we're reading from + current_file: Option<SpillPoolFile>, + /// Schema of the spilled data + schema: SchemaRef, +} + +impl SpillPoolReader { + /// Creates a new reader from shared pool state. + /// + /// This is private - use the [`spsc_channel`] function to create a reader/writer pair. + /// + /// # Arguments + /// + /// * `shared` - Shared reference to the pool state + fn new(shared: Arc<Mutex<SpillPoolShared>>, schema: SchemaRef) -> Self { + Self { + shared, + current_file: None, + schema, + } + } +} + +impl Stream for SpillPoolReader { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll<Option<Self::Item>> { + use std::task::Poll; + + loop { + // If we have a current file, try to read from it + if let Some(ref mut file) = self.current_file { + match file.poll_next_unpin(cx) { + Poll::Ready(Some(Ok(batch))) => { + // Got a batch, return it + return Poll::Ready(Some(Ok(batch))); + } + Poll::Ready(Some(Err(e))) => { + // Error reading batch + return Poll::Ready(Some(Err(e))); + } + Poll::Ready(None) => { + // Current file stream exhausted + // Check if this file is marked as writer_finished + let writer_finished = { file.shared.lock().writer_finished }; + + if writer_finished { + // File is complete, pop it from the queue and move to next + let mut shared = self.shared.lock(); + shared.files.pop_front(); + drop(shared); // Release lock + + // Clear current file and continue loop to get next file + self.current_file = None; + continue; + } else { + // Stream exhausted but writer not finished - unexpected + // This shouldn't happen with proper coordination + return Poll::Ready(None); + } + } + Poll::Pending => { + // File not ready yet (waiting for writer) + // Register waker so we get notified when writer adds more batches + let mut shared = self.shared.lock(); + shared.register_waker(cx.waker().clone()); + return Poll::Pending; + } + } + } + + // No current file, need to get the next one + let mut shared = self.shared.lock(); + + // Peek at the front of the queue (don't pop yet) + if let Some(file_shared) = shared.files.front() { + // Create a SpillPoolFile from the shared state + let spill_manager = Arc::clone(&shared.spill_manager); + let file_shared = Arc::clone(file_shared); + drop(shared); // Release lock before creating SpillPoolFile + + self.current_file = Some(SpillPoolFile { + shared: file_shared, + reader: None, + spill_manager, + }); + + // Continue loop to poll the new file + continue; + } + + // No files in queue - check if writer is done + if shared.remaining_writer_count == 0 { + // Writer is done and no more files will be added - EOF + return Poll::Ready(None); + } + + // Writer still active, register waker that will get notified when new files are added + shared.register_waker(cx.waker().clone()); + return Poll::Pending; + } + } +} + +impl RecordBatchStream for SpillPoolReader { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; + use arrow::array::{ArrayRef, Int32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common_runtime::{JoinSet, SpawnedTask}; + use datafusion_execution::runtime_env::RuntimeEnv; + + fn create_test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn create_test_batch(start: i32, count: usize) -> RecordBatch { + let schema = create_test_schema(); + let a: ArrayRef = Arc::new(Int32Array::from( + (start..start + count as i32).collect::<Vec<_>>(), + )); + RecordBatch::try_new(schema, vec![a]).unwrap() + } + + fn create_spill_channel( + max_file_size: usize, + ) -> (SpillPoolSink, SendableRecordBatchStream) { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); + + spsc_channel(max_file_size, spill_manager) + } + + fn create_shared_spill_channel( + max_file_size: usize, + ) -> (SpillPoolWriter, SendableRecordBatchStream) { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); + + mpsc_channel(max_file_size, spill_manager) + } + + fn create_spill_channel_with_metrics( + max_file_size: usize, + ) -> (SpillPoolSink, SendableRecordBatchStream, SpillMetrics) { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = Arc::new(SpillManager::new(env, metrics.clone(), schema)); + + let (writer, reader) = spsc_channel(max_file_size, spill_manager); + (writer, reader, metrics) + } + + #[tokio::test] + async fn test_basic_write_and_read() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + // Write one batch + let batch1 = create_test_batch(0, 10); + writer.push_batch(&batch1)?; + + // Read the batch + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 10); + + // Write another batch + let batch2 = create_test_batch(10, 5); + writer.push_batch(&batch2)?; + // Read the second batch + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 5); + + Ok(()) + } + + #[tokio::test] + async fn test_single_batch_write_read() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + // Write one batch + let batch = create_test_batch(0, 5); + writer.push_batch(&batch)?; + + // Read it back + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 5); + + // Verify the actual data + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), 0); + assert_eq!(col.value(4), 4); + + Ok(()) + } + + #[tokio::test] + async fn test_multiple_batches_sequential() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + // Write multiple batches + for i in 0..5 { + let batch = create_test_batch(i * 10, 10); + writer.push_batch(&batch)?; + } + + // Read all batches and verify FIFO order + for i in 0..5 { + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 10); + + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), i * 10, "Batch {i} not in FIFO order"); + } + + Ok(()) + } + + #[tokio::test] + async fn test_empty_writer() -> Result<()> { + let (_writer, reader) = create_spill_channel(1024 * 1024); + + // Reader should pend since no batches were written + let mut reader = reader; + let result = + tokio::time::timeout(std::time::Duration::from_millis(100), reader.next()) + .await; + + assert!(result.is_err(), "Reader should timeout on empty writer"); + + Ok(()) + } + + #[tokio::test] + async fn test_empty_batch_skipping() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + // Write empty batch + let empty_batch = create_test_batch(0, 0); + writer.push_batch(&empty_batch)?; + + // Write non-empty batch + let batch = create_test_batch(0, 5); + writer.push_batch(&batch)?; + + // Should only read the non-empty batch + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 5); + + Ok(()) + } + + #[tokio::test] + async fn test_rotation_triggered_by_size() -> Result<()> { + // Set a small max_file_size to trigger rotation after one batch + let batch1 = create_test_batch(0, 10); + let batch_size = batch1.get_array_memory_size() + 1; + + let (writer, mut reader, metrics) = create_spill_channel_with_metrics(batch_size); + + // Write first batch (should fit in first file) + writer.push_batch(&batch1)?; + + // Check metrics after first batch - file created but not finalized yet + assert_eq!( + metrics.spill_file_count.value(), + 1, + "Should have created 1 file after first batch" + ); + assert_eq!( + metrics.spilled_bytes.value(), + 320, + "Spilled bytes should reflect data written (header + 1 batch)" + ); + assert_eq!( + metrics.spilled_rows.value(), + 10, + "Should have spilled 10 rows from first batch" + ); + + // Write second batch (should trigger rotation - finalize first file) + let batch2 = create_test_batch(10, 10); + assert!( + batch2.get_array_memory_size() <= batch_size, + "batch2 size {} exceeds limit {batch_size}", + batch2.get_array_memory_size(), + ); + assert!( + batch1.get_array_memory_size() + batch2.get_array_memory_size() > batch_size, + "Combined size {} does not exceed limit to trigger rotation", + batch1.get_array_memory_size() + batch2.get_array_memory_size() + ); + writer.push_batch(&batch2)?; + + // Check metrics after rotation - first file finalized, but second file not created yet + // (new file created lazily on next push_batch call) + assert_eq!( + metrics.spill_file_count.value(), + 1, + "Should still have 1 file (second file not created until next write)" + ); + assert!( + metrics.spilled_bytes.value() > 0, + "Spilled bytes should be > 0 after first file finalized (got {})", + metrics.spilled_bytes.value() + ); + assert_eq!( + metrics.spilled_rows.value(), + 20, + "Should have spilled 20 total rows (10 + 10)" + ); + + // Write a third batch to confirm rotation occurred (creates second file) + let batch3 = create_test_batch(20, 5); + writer.push_batch(&batch3)?; + + // Now check that second file was created + assert_eq!( + metrics.spill_file_count.value(), + 2, + "Should have created 2 files after writing to new file" + ); + assert_eq!( + metrics.spilled_rows.value(), + 25, + "Should have spilled 25 total rows (10 + 10 + 5)" + ); + + // Read all three batches + let result1 = reader.next().await.unwrap()?; + assert_eq!(result1.num_rows(), 10); + + let result2 = reader.next().await.unwrap()?; + assert_eq!(result2.num_rows(), 10); + + let result3 = reader.next().await.unwrap()?; + assert_eq!(result3.num_rows(), 5); + + Ok(()) + } + + #[tokio::test] + async fn test_multiple_rotations() -> Result<()> { + let batches = (0..10) + .map(|i| create_test_batch(i * 10, 10)) + .collect::<Vec<_>>(); + + let batch_size = batches[0].get_array_memory_size() * 2 + 1; + + // Very small max_file_size to force frequent rotations + let (writer, mut reader, metrics) = create_spill_channel_with_metrics(batch_size); + + // Write many batches to cause multiple rotations + for i in 0..10 { + let batch = create_test_batch(i * 10, 10); + writer.push_batch(&batch)?; + } + + // Check metrics after all writes - should have multiple files due to rotations + // With batch_size = 2 * one_batch + 1, each file fits ~2 batches before rotating + // 10 batches should create multiple files (exact count depends on rotation timing) + let file_count = metrics.spill_file_count.value(); + assert!( + file_count >= 4, + "Should have created at least 4 files with multiple rotations (got {file_count})" + ); + assert!( + metrics.spilled_bytes.value() > 0, + "Spilled bytes should be > 0 after rotations (got {})", + metrics.spilled_bytes.value() + ); + assert_eq!( + metrics.spilled_rows.value(), + 100, + "Should have spilled 100 total rows (10 batches * 10 rows)" + ); + + // Read all batches and verify order + for i in 0..10 { + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 10); + + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!( + col.value(0), + i * 10, + "Batch {i} not in correct order after rotations" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_single_batch_larger_than_limit() -> Result<()> { + // Very small limit + let (writer, mut reader, metrics) = create_spill_channel_with_metrics(100); + + // Write a batch that exceeds the limit + let large_batch = create_test_batch(0, 100); + writer.push_batch(&large_batch)?; + + // Check metrics after large batch - should trigger rotation immediately + assert_eq!( + metrics.spill_file_count.value(), + 1, + "Should have created 1 file for large batch" + ); + assert_eq!( + metrics.spilled_rows.value(), + 100, + "Should have spilled 100 rows from large batch" + ); + + // Should still write and read successfully + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 100); + + // Next batch should go to a new file + let batch2 = create_test_batch(100, 10); + writer.push_batch(&batch2)?; + + // Check metrics after second batch - should have rotated to a new file + assert_eq!( + metrics.spill_file_count.value(), + 2, + "Should have created 2 files after rotation" + ); + assert_eq!( + metrics.spilled_rows.value(), + 110, + "Should have spilled 110 total rows (100 + 10)" + ); + + let result2 = reader.next().await.unwrap()?; + assert_eq!(result2.num_rows(), 10); + + Ok(()) + } + + #[tokio::test] + async fn test_very_small_max_file_size() -> Result<()> { + // Test with just 1 byte max (extreme case) + let (writer, mut reader) = create_spill_channel(1); + + // Any batch will exceed this limit + let batch = create_test_batch(0, 5); + writer.push_batch(&batch)?; + + // Should still work + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 5); + + Ok(()) + } + + #[tokio::test] + async fn test_exact_size_boundary() -> Result<()> { + // Create a batch and measure its approximate size + let batch = create_test_batch(0, 10); + let batch_size = batch.get_array_memory_size(); + + // Set max_file_size to exactly the batch size + let (writer, mut reader, metrics) = create_spill_channel_with_metrics(batch_size); + + // Write first batch (exactly at the size limit) + writer.push_batch(&batch)?; + + // Check metrics after first batch - should NOT rotate yet (size == limit, not >) + assert_eq!( + metrics.spill_file_count.value(), + 1, + "Should have created 1 file after first batch at exact boundary" + ); + assert_eq!( + metrics.spilled_rows.value(), + 10, + "Should have spilled 10 rows from first batch" + ); + + // Write second batch (exceeds the limit, should trigger rotation) + let batch2 = create_test_batch(10, 10); + writer.push_batch(&batch2)?; + + // Check metrics after second batch - rotation triggered, first file finalized + // Note: second file not created yet (lazy creation on next write) + assert_eq!( + metrics.spill_file_count.value(), + 1, + "Should still have 1 file after rotation (second file created lazily)" + ); + assert_eq!( + metrics.spilled_rows.value(), + 20, + "Should have spilled 20 total rows (10 + 10)" + ); + // Verify first file was finalized by checking spilled_bytes + assert!( + metrics.spilled_bytes.value() > 0, + "Spilled bytes should be > 0 after file finalization (got {})", + metrics.spilled_bytes.value() + ); + + // Both should be readable + let result1 = reader.next().await.unwrap()?; + assert_eq!(result1.num_rows(), 10); + + let result2 = reader.next().await.unwrap()?; + assert_eq!(result2.num_rows(), 10); + + // Spill another batch, now we should see the second file created + let batch3 = create_test_batch(20, 5); + writer.push_batch(&batch3)?; + assert_eq!( + metrics.spill_file_count.value(), + 2, + "Should have created 2 files after writing to new file" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_concurrent_reader_writer() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + // Spawn writer task + let writer_handle = SpawnedTask::spawn(async move { + for i in 0..10 { + let batch = create_test_batch(i * 10, 10); + writer.push_batch(&batch).unwrap(); + // Small delay to simulate real concurrent work + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + }); + + // Reader task (runs concurrently) + let reader_handle = SpawnedTask::spawn(async move { + let mut count = 0; + for i in 0..10 { + let result = reader.next().await.unwrap().unwrap(); + assert_eq!(result.num_rows(), 10); + + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), i * 10); + count += 1; + } + count + }); + + // Wait for both to complete + writer_handle.await.unwrap(); + let batches_read = reader_handle.await.unwrap(); + assert_eq!(batches_read, 10); + + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 10)] + async fn test_concurrent_writers() -> Result<()> { + let (writer, mut reader) = create_shared_spill_channel(1024 * 1024); + + // Spawn writer tasks + let mut writer_join_set = JoinSet::new(); + for w in 0..10 { + let writer = writer.clone(); + writer_join_set.spawn(async move { + for b in 0..10 { + let batch = create_test_batch((w * 100) + (b * 10), 10); + writer.push_batch(&batch).unwrap(); + } + }); + } + drop(writer); + + // Reader task (runs concurrently) + let reader_handle = SpawnedTask::spawn(async move { + let mut batch_order = vec![]; + loop { + match reader.next().await { + None => break, + Some(batch) => { + let batch = batch.unwrap(); + + assert_eq!(batch.num_rows(), 10); + + let col = batch + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + batch_order.push(col.value(0) / 10); + } + } + } + batch_order + }); + + // Wait for both to complete + writer_join_set.join_all().await; + let mut batch_order = reader_handle.await.unwrap(); + + // When used with multiple writers, order is not guaranteed + batch_order.sort(); + assert_eq!(batch_order, (0i32..100i32).collect::<Vec<_>>()); + + Ok(()) + } + + #[tokio::test] + async fn test_reader_catches_up_to_writer() -> Result<()> { + let (writer, mut reader) = create_spill_channel(1024 * 1024); + + let (reader_waiting_tx, reader_waiting_rx) = tokio::sync::oneshot::channel(); + let (first_read_done_tx, first_read_done_rx) = tokio::sync::oneshot::channel(); + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum ReadWriteEvent { + ReadStart, + Read(usize), + Write(usize), + } + + let events = Arc::new(Mutex::new(vec![])); + // Start reader first (will pend) + let reader_events = Arc::clone(&events); + let reader_handle = SpawnedTask::spawn(async move { + reader_events.lock().push(ReadWriteEvent::ReadStart); + reader_waiting_tx + .send(()) + .expect("reader_waiting channel closed unexpectedly"); + let result = reader.next().await.unwrap().unwrap(); + reader_events + .lock() + .push(ReadWriteEvent::Read(result.num_rows())); + first_read_done_tx + .send(()) + .expect("first_read_done channel closed unexpectedly"); + let result = reader.next().await.unwrap().unwrap(); + reader_events + .lock() + .push(ReadWriteEvent::Read(result.num_rows())); + }); + + // Wait until the reader is pending on the first batch + reader_waiting_rx + .await + .expect("reader should signal when waiting"); + + // Now write a batch (should wake the reader) + let batch = create_test_batch(0, 5); + events.lock().push(ReadWriteEvent::Write(batch.num_rows())); + writer.push_batch(&batch)?; + + // Wait for the reader to finish the first read before allowing the + // second write. This ensures deterministic ordering of events: + // 1. The reader starts and pends on the first `next()` + // 2. The first write wakes the reader + // 3. The reader processes the first batch and signals completion + // 4. The second write is issued, ensuring consistent event ordering + first_read_done_rx + .await + .expect("reader should signal when first read completes"); + + // Write another batch + let batch = create_test_batch(5, 10); + events.lock().push(ReadWriteEvent::Write(batch.num_rows())); + writer.push_batch(&batch)?; + + // Reader should complete + reader_handle.await.unwrap(); + let events = events.lock().clone(); + assert_eq!( + events, + vec![ + ReadWriteEvent::ReadStart, + ReadWriteEvent::Write(5), + ReadWriteEvent::Read(5), + ReadWriteEvent::Write(10), + ReadWriteEvent::Read(10) + ] + ); + + Ok(()) + } + + #[tokio::test] + async fn test_reader_starts_after_writer_finishes() -> Result<()> { + let (writer, reader) = create_spill_channel(128); + + // Writer writes all data + for i in 0..5 { + let batch = create_test_batch(i * 10, 10); + writer.push_batch(&batch)?; + } + + drop(writer); + + // Now start reader + let mut reader = reader; + let mut count = 0; + for i in 0..5 { + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 10); + + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), i * 10); + count += 1; + } + + assert_eq!(count, 5, "Should read all batches after writer finishes"); + + Ok(()) + } + + #[tokio::test] + async fn test_writer_drop_finalizes_file() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = + Arc::new(SpillManager::new(Arc::clone(&env), metrics.clone(), schema)); + + let (writer, mut reader) = spsc_channel(1024 * 1024, spill_manager); + + // Write some batches + for i in 0..5 { + let batch = create_test_batch(i * 10, 10); + writer.push_batch(&batch)?; + } + + // Check metrics before drop - spilled_bytes already reflects written data + let spilled_bytes_before = metrics.spilled_bytes.value(); + assert_eq!( + spilled_bytes_before, 1088, + "Spilled bytes should reflect data written (header + 5 batches)" + ); + + // Explicitly drop the writer - this should finalize the current file + drop(writer); + + // Check metrics after drop - spilled_bytes should be > 0 now + let spilled_bytes_after = metrics.spilled_bytes.value(); + assert!( + spilled_bytes_after > 0, + "Spilled bytes should be > 0 after writer is dropped (got {spilled_bytes_after})" + ); + + // Verify reader can still read all batches + let mut count = 0; + for i in 0..5 { + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), 10); + + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), i * 10); + count += 1; + } + + assert_eq!(count, 5, "Should read all batches after writer is dropped"); + + Ok(()) + } + + /// Verifies that the reader stays alive as long as any writer clone exists. + /// + /// `SpillPoolWriter` is `Clone`, and in non-preserve-order repartitioning + /// mode multiple input partition tasks share clones of the same writer. + /// The reader must not see EOF until **all** clones have been dropped, + /// even if the queue is temporarily empty between writes from different + /// clones. + /// + /// The test sequence is: + /// + /// 1. writer1 writes a batch, then is dropped. + /// 2. The reader consumes that batch (queue is now empty). + /// 3. writer2 (still alive) writes a batch. + /// 4. The reader must see that batch. + /// 5. EOF is only signalled after writer2 is also dropped. + #[tokio::test] + async fn test_clone_drop_does_not_signal_eof_prematurely() -> Result<()> { + let (writer1, mut reader) = create_shared_spill_channel(1024 * 1024); + let writer2 = writer1.clone(); + + // Synchronization: tell writer2 when it may proceed. + let (proceed_tx, proceed_rx) = tokio::sync::oneshot::channel::<()>(); + + // Spawn writer2 — it waits for the signal before writing. + let writer2_handle = SpawnedTask::spawn(async move { + proceed_rx.await.unwrap(); + writer2.push_batch(&create_test_batch(10, 10)).unwrap(); + // writer2 is dropped here (last clone → true EOF) + }); + + // Writer1 writes one batch, then drops. + writer1.push_batch(&create_test_batch(0, 10))?; + drop(writer1); + + // Read writer1's batch. + let batch1 = reader.next().await.unwrap()?; + assert_eq!(batch1.num_rows(), 10); + let col = batch1 + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), 0); + + // Signal writer2 to write its batch. It will execute when the + // current task yields (i.e. when reader.next() returns Pending). + proceed_tx.send(()).unwrap(); + + // The reader should wait (Pending) for writer2's data, not EOF. + let batch2 = + tokio::time::timeout(std::time::Duration::from_secs(5), reader.next()) + .await + .expect("Reader timed out — should not hang"); + + assert!( + batch2.is_some(), + "Reader must not return EOF while a writer clone is still alive" + ); + let batch2 = batch2.unwrap()?; + assert_eq!(batch2.num_rows(), 10); + let col = batch2 + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), 10); + + writer2_handle.await.unwrap(); + + // All writers dropped — reader should see real EOF now. + assert!(reader.next().await.is_none()); + + Ok(()) + } + + #[tokio::test] + async fn test_disk_usage_decreases_as_files_consumed() -> Result<()> { + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + // Test configuration + const NUM_BATCHES: usize = 3; + const ROWS_PER_BATCH: usize = 100; + + // Step 1: Create a test batch and measure its size + let batch = create_test_batch(0, ROWS_PER_BATCH); + let batch_size = batch.get_array_memory_size(); + + // Step 2: Configure file rotation to approximately 1 batch per file + // Create a custom RuntimeEnv so we can access the DiskManager + let runtime = Arc::new(RuntimeEnvBuilder::default().build()?); + let disk_manager = Arc::clone(&runtime.disk_manager); + + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = Arc::new(SpillManager::new(runtime, metrics.clone(), schema)); + + let (writer, mut reader) = spsc_channel(batch_size - 1, spill_manager); + + // Step 3: Write NUM_BATCHES batches to create approximately NUM_BATCHES files + for i in 0..NUM_BATCHES { + let start = (i * ROWS_PER_BATCH) as i32; + writer.push_batch(&create_test_batch(start, ROWS_PER_BATCH))?; + } + + // Check how many files were created (should be at least a few due to file rotation) + let file_count = metrics.spill_file_count.value(); + assert_eq!( + file_count, NUM_BATCHES, + "Expected at {NUM_BATCHES} files with rotation, got {file_count}" + ); + + // Step 4: Verify initial disk usage reflects all files + let initial_disk_usage = disk_manager.used_disk_space(); + assert!( + initial_disk_usage > 0, + "Expected disk usage > 0 after writing batches, got {initial_disk_usage}" + ); + + // Step 5: Read NUM_BATCHES - 1 batches (all but 1) + // As each file is fully consumed, it should be dropped and disk usage should decrease + for i in 0..(NUM_BATCHES - 1) { + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), ROWS_PER_BATCH); + + let col = result + .column(0) + .as_any() + .downcast_ref::<Int32Array>() + .unwrap(); + assert_eq!(col.value(0), (i * ROWS_PER_BATCH) as i32); + } + + // Step 6: Verify disk usage decreased but is not zero (at least 1 batch remains) + let partial_disk_usage = disk_manager.used_disk_space(); + assert!( + partial_disk_usage > 0 + && partial_disk_usage < (batch_size * NUM_BATCHES * 2) as u64, + "Disk usage should be > 0 with remaining batches" + ); + assert!( + partial_disk_usage < initial_disk_usage, + "Disk usage should have decreased after reading most batches: initial={initial_disk_usage}, partial={partial_disk_usage}" + ); + + // Step 7: Read the final batch + let result = reader.next().await.unwrap()?; + assert_eq!(result.num_rows(), ROWS_PER_BATCH); + + // Step 8: Drop writer first to signal no more data will be written + // The reader has infinite stream semantics and will wait for the writer + // to be dropped before returning None + drop(writer); + + // Verify we've read all batches - now the reader should return None + assert!( + reader.next().await.is_none(), + "Should have no more batches to read" + ); + + // Step 9: Drop reader to release all references + drop(reader); + + // Step 10: Verify complete cleanup - disk usage should be 0 + let final_disk_usage = disk_manager.used_disk_space(); + assert_eq!( + final_disk_usage, 0, + "Disk usage should be 0 after all files dropped, got {final_disk_usage}" + ); + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/statistics.rs b/native/vendor/datafusion-physical-plan/src/statistics.rs new file mode 100644 index 0000000..9246d7d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/statistics.rs
@@ -0,0 +1,274 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Statistics computation for physical plans. +//! +//! [`StatisticsArgs`] provides external context to +//! [`ExecutionPlan::statistics_from_inputs`]. + +use crate::ExecutionPlan; +use datafusion_common::{ + Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, +}; +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; +use std::sync::Arc; + +/// Per-call memoization cache for statistics computation. +/// +/// Keyed by `(plan node pointer address, partition)`. Shared across +/// a single statistics walk via [`StatisticsContext`]. +/// +/// The pointer-based key is safe within a single synchronous walk: +/// all `Arc<dyn ExecutionPlan>` nodes are held by the plan tree for +/// the duration of the walk, so addresses cannot be reused. +#[derive(Debug, Default)] +struct StatsCache(HashMap<(usize, Option<usize>), Arc<Statistics>>); + +impl StatsCache { + fn get( + &self, + plan: &dyn ExecutionPlan, + partition: Option<usize>, + ) -> Option<&Arc<Statistics>> { + let key = ( + plan as *const dyn ExecutionPlan as *const () as usize, + partition, + ); + self.0.get(&key) + } + + fn insert( + &mut self, + plan: &dyn ExecutionPlan, + partition: Option<usize>, + stats: Arc<Statistics>, + ) { + let key = ( + plan as *const dyn ExecutionPlan as *const () as usize, + partition, + ); + self.0.insert(key, stats); + } +} + +/// Arguments passed to [`ExecutionPlan::statistics_from_inputs`] carrying +/// external information that operators can use when computing their +/// statistics. +#[derive(Debug, Default, Clone)] +pub struct StatisticsArgs { + partition: Option<usize>, +} + +impl StatisticsArgs { + /// Creates new statistics arguments. + /// + /// By default the partition is set to `None` (statistics should be computed + /// for the entire plan). + pub fn new() -> Self { + Default::default() + } + + /// Set the partition to compute statistics + /// + /// * `None` means statistics should be computed for the entire plan. + /// * `Some(idx)` means statistics should be computed for the specified + /// partition index. + pub fn set_partition(&mut self, partition: Option<usize>) { + self.partition = partition; + } + + /// Builder Style API for [`Self::set_partition`] + pub fn with_partition(mut self, partition: Option<usize>) -> Self { + self.set_partition(partition); + self + } + + /// Return the partition to compute statistics + pub fn partition(&self) -> Option<usize> { + self.partition + } +} + +/// Directive returned by [`ExecutionPlan::child_stats_requests`] describing +/// how the [`StatisticsContext`] should obtain each child's statistics. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildStats { + /// Compute the child's statistics at this partition (`None` = overall). + At(Option<usize>), + /// Skip this child; the parent does not need its statistics. A placeholder + /// [`Statistics::new_unknown`] is supplied in its slot. + Skip, +} + +/// Owns the bottom-up traversal and per-walk memoization cache for statistics +/// computation. Call [`StatisticsContext::compute`] to walk a plan tree. +pub struct StatisticsContext { + cache: Rc<RefCell<StatsCache>>, +} + +impl Default for StatisticsContext { + fn default() -> Self { + Self::new() + } +} + +impl StatisticsContext { + /// Creates a context with an empty cache. + pub fn new() -> Self { + Self { + cache: Rc::new(RefCell::new(StatsCache::default())), + } + } + + /// Clears the memoization cache. + /// + /// The cache is keyed by raw plan-node pointers, which are only stable + /// while the current plan tree is alive. Reset between optimizer passes + /// (which rewrite the plan) when reusing one context across them, so stale + /// pointer keys cannot collide. + pub fn reset_cache(&self) { + self.cache.borrow_mut().0.clear(); + } + + /// Computes statistics for `plan`, resolving children first and passing + /// the results to [`ExecutionPlan::statistics_from_inputs`]. + /// + /// When `args.partition()` is `Some(idx)`, `idx` is validated against the + /// plan's partition count. + pub fn compute( + &self, + plan: &dyn ExecutionPlan, + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let partition = args.partition(); + + if let Some(idx) = partition { + let partition_count = plan.properties().partitioning.partition_count(); + assert_or_internal_err!( + idx < partition_count, + "Invalid partition index: {}, the partition count is {}", + idx, + partition_count + ); + } + + if let Some(cached) = self.cache.borrow().get(plan, partition) { + return Ok(Arc::clone(cached)); + } + + let children = plan.children(); + let requests = plan.child_stats_requests(partition); + assert_eq_or_internal_err!( + requests.len(), + children.len(), + "{} child_stats_requests returned {} entries for {} children", + plan.name(), + requests.len(), + children.len() + ); + let child_stats = children + .iter() + .zip(requests) + .map(|(child, directive)| match directive { + ChildStats::At(p) => { + self.compute(child.as_ref(), &StatisticsArgs::new().with_partition(p)) + } + ChildStats::Skip => { + Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) + } + }) + .collect::<Result<Vec<_>>>()?; + + let result = plan.statistics_from_inputs(&child_stats, args)?; + self.cache + .borrow_mut() + .insert(plan, partition, Arc::clone(&result)); + Ok(result) + } +} + +#[cfg(all(test, feature = "test_utils"))] +mod tests { + use super::*; + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::test::exec::StatisticsExec; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ColumnStatistics, stats::Precision}; + + fn make_stats_leaf(num_rows: usize) -> Arc<dyn ExecutionPlan> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let col_stats = vec![ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + byte_size: Precision::Absent, + }]; + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(num_rows), + total_byte_size: Precision::Absent, + column_statistics: col_stats, + }, + schema, + )) + } + + #[test] + fn coalesce_returns_overall_stats_for_any_partition() { + let leaf = make_stats_leaf(100); + let plan: Arc<dyn ExecutionPlan> = Arc::new(CoalescePartitionsExec::new(leaf)); + + let ctx = StatisticsContext::new(); + let stats = ctx + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(100)); + + let stats_none = ctx.compute(plan.as_ref(), &StatisticsArgs::new()).unwrap(); + assert_eq!(stats_none.num_rows, Precision::Exact(100)); + } + + #[test] + fn context_caches_within_walk() { + let leaf = make_stats_leaf(42); + let ctx = StatisticsContext::new(); + let args = StatisticsArgs::new(); + + let s1 = ctx.compute(leaf.as_ref(), &args).unwrap(); + assert!(!ctx.cache.borrow().0.is_empty()); + + let s2 = ctx.compute(leaf.as_ref(), &args).unwrap(); + assert!(Arc::ptr_eq(&s1, &s2)); + } + + #[test] + fn reset_cache_clears_entries() { + let leaf = make_stats_leaf(10); + let ctx = StatisticsContext::new(); + let _ = ctx.compute(leaf.as_ref(), &StatisticsArgs::new()).unwrap(); + assert!(!ctx.cache.borrow().0.is_empty()); + ctx.reset_cache(); + assert!(ctx.cache.borrow().0.is_empty()); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/stream.rs b/native/vendor/datafusion-physical-plan/src/stream.rs new file mode 100644 index 0000000..9d0b964 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/stream.rs
@@ -0,0 +1,1163 @@ +// 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. + +//! Stream wrappers for physical operators + +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +#[cfg(test)] +use super::metrics::ExecutionPlanMetricsSet; +use super::metrics::{BaselineMetrics, SplitMetrics}; +use super::{ExecutionPlan, RecordBatchStream, SendableRecordBatchStream}; +use crate::displayable; +use crate::spill::get_record_batch_memory_size; + +use arrow::{datatypes::SchemaRef, record_batch::RecordBatch}; +use datafusion_common::{Result, exec_err}; +use datafusion_common_runtime::JoinSet; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryReservation; + +use futures::ready; +use futures::stream::BoxStream; +use futures::{Future, Stream, StreamExt}; +use log::debug; +use pin_project_lite::pin_project; +use tokio::runtime::Handle; +use tokio::sync::mpsc::{Receiver, Sender}; + +/// Creates a stream from a collection of producing tasks, routing panics to the stream. +/// +/// Note that this is similar to [`ReceiverStream` from tokio-stream], with the differences being: +/// +/// 1. Methods to bound and "detach" tasks (`spawn()` and `spawn_blocking()`). +/// +/// 2. Propagates panics, whereas the `tokio` version doesn't propagate panics to the receiver. +/// +/// 3. Automatically cancels any outstanding tasks when the receiver stream is dropped. +/// +/// [`ReceiverStream` from tokio-stream]: https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.ReceiverStream.html +pub(crate) struct ReceiverStreamBuilder<O> { + tx: Sender<Result<O>>, + rx: Receiver<Result<O>>, + join_set: JoinSet<Result<()>>, +} + +impl<O: Send + 'static> ReceiverStreamBuilder<O> { + /// Create new channels with the specified buffer size + pub fn new(capacity: usize) -> Self { + let (tx, rx) = tokio::sync::mpsc::channel(capacity); + + Self { + tx, + rx, + join_set: JoinSet::new(), + } + } + + /// Get a handle for sending data to the output + pub fn tx(&self) -> Sender<Result<O>> { + self.tx.clone() + } + + /// Spawn task that will be aborted if this builder (or the stream + /// built from it) are dropped + pub fn spawn<F>(&mut self, task: F) + where + F: Future<Output = Result<()>>, + F: Send + 'static, + { + self.join_set.spawn(task); + } + + /// Same as [`Self::spawn`] but it spawns the task on the provided runtime + pub fn spawn_on<F>(&mut self, task: F, handle: &Handle) + where + F: Future<Output = Result<()>>, + F: Send + 'static, + { + self.join_set.spawn_on(task, handle); + } + + /// Spawn a blocking task that will be aborted if this builder (or the stream + /// built from it) are dropped. + /// + /// This is often used to spawn tasks that write to the sender + /// retrieved from `Self::tx`. + pub fn spawn_blocking<F>(&mut self, f: F) + where + F: FnOnce() -> Result<()>, + F: Send + 'static, + { + self.join_set.spawn_blocking(f); + } + + /// Same as [`Self::spawn_blocking`] but it spawns the blocking task on the provided runtime + pub fn spawn_blocking_on<F>(&mut self, f: F, handle: &Handle) + where + F: FnOnce() -> Result<()>, + F: Send + 'static, + { + self.join_set.spawn_blocking_on(f, handle); + } + + /// Create a stream of all data written to `tx` + pub fn build(self) -> BoxStream<'static, Result<O>> { + let Self { + tx, + rx, + mut join_set, + } = self; + + // Doesn't need tx + drop(tx); + + // future that checks the result of the join set, and propagates panic if seen + let check = async move { + while let Some(result) = join_set.join_next().await { + match result { + Ok(task_result) => { + match task_result { + // Nothing to report + Ok(_) => continue, + // This means a blocking task error + Err(error) => return Some(Err(error)), + } + } + // This means a tokio task error, likely a panic + Err(e) => { + if e.is_panic() { + // resume on the main thread + std::panic::resume_unwind(e.into_panic()); + } else { + // This should only occur if the task is + // cancelled, which would only occur if + // the JoinSet were aborted, which in turn + // would imply that the receiver has been + // dropped and this code is not running + return Some(exec_err!("Non Panic Task error: {e}")); + } + } + } + } + None + }; + + let check_stream = futures::stream::once(check) + // unwrap Option / only return the error + .filter_map(|item| async move { item }); + + // Convert the receiver into a stream + let rx_stream = futures::stream::unfold(rx, |mut rx| async move { + let next_item = rx.recv().await; + next_item.map(|next_item| (next_item, rx)) + }); + + // Merge the streams together so whichever is ready first + // produces the batch + futures::stream::select(rx_stream, check_stream).boxed() + } +} + +/// Builder for `RecordBatchReceiverStream` that propagates errors +/// and panic's correctly. +/// +/// [`RecordBatchReceiverStreamBuilder`] is used to spawn one or more tasks +/// that produce [`RecordBatch`]es and send them to a single +/// `Receiver` which can improve parallelism. +/// +/// This also handles propagating panic`s and canceling the tasks. +/// +/// # Example +/// +/// The following example spawns 2 tasks that will write [`RecordBatch`]es to +/// the `tx` end of the builder, after building the stream, we can receive +/// those batches with calling `.next()` +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion_common::arrow::datatypes::{Schema, Field, DataType}; +/// # use datafusion_common::arrow::array::RecordBatch; +/// # use datafusion_physical_plan::stream::RecordBatchReceiverStreamBuilder; +/// # use futures::stream::StreamExt; +/// # use tokio::runtime::Builder; +/// # let rt = Builder::new_current_thread().build().unwrap(); +/// # +/// # rt.block_on(async { +/// let schema = Arc::new(Schema::new(vec![Field::new("foo", DataType::Int8, false)])); +/// let mut builder = RecordBatchReceiverStreamBuilder::new(Arc::clone(&schema), 10); +/// +/// // task 1 +/// let tx_1 = builder.tx(); +/// let schema_1 = Arc::clone(&schema); +/// builder.spawn(async move { +/// // Your task needs to send batches to the tx +/// tx_1.send(Ok(RecordBatch::new_empty(schema_1))) +/// .await +/// .unwrap(); +/// +/// Ok(()) +/// }); +/// +/// // task 2 +/// let tx_2 = builder.tx(); +/// let schema_2 = Arc::clone(&schema); +/// builder.spawn(async move { +/// // Your task needs to send batches to the tx +/// tx_2.send(Ok(RecordBatch::new_empty(schema_2))) +/// .await +/// .unwrap(); +/// +/// Ok(()) +/// }); +/// +/// let mut stream = builder.build(); +/// while let Some(res_batch) = stream.next().await { +/// // `res_batch` can either from task 1 or 2 +/// +/// // do something with `res_batch` +/// } +/// # }); +/// ``` +pub struct RecordBatchReceiverStreamBuilder { + schema: SchemaRef, + inner: ReceiverStreamBuilder<RecordBatch>, +} + +impl RecordBatchReceiverStreamBuilder { + /// Create new channels with the specified buffer size + pub fn new(schema: SchemaRef, capacity: usize) -> Self { + Self { + schema, + inner: ReceiverStreamBuilder::new(capacity), + } + } + + /// Get a handle for sending [`RecordBatch`] to the output + /// + /// If the stream is dropped / canceled, the sender will be closed and + /// calling `tx().send()` will return an error. Producers should stop + /// producing in this case and return control. + pub fn tx(&self) -> Sender<Result<RecordBatch>> { + self.inner.tx() + } + + /// Spawn task that will be aborted if this builder (or the stream + /// built from it) are dropped + /// + /// This is often used to spawn tasks that write to the sender + /// retrieved from [`Self::tx`], for examples, see the document + /// of this type. + pub fn spawn<F>(&mut self, task: F) + where + F: Future<Output = Result<()>>, + F: Send + 'static, + { + self.inner.spawn(task) + } + + /// Same as [`Self::spawn`] but it spawns the task on the provided runtime. + pub fn spawn_on<F>(&mut self, task: F, handle: &Handle) + where + F: Future<Output = Result<()>>, + F: Send + 'static, + { + self.inner.spawn_on(task, handle) + } + + /// Spawn a blocking task tied to the builder and stream. + /// + /// # Drop / Cancel Behavior + /// + /// If this builder (or the stream built from it) is dropped **before** the + /// task starts, the task is also dropped and will never start execute. + /// + /// **Note:** Once the blocking task has started, it **will not** be + /// forcibly stopped on drop as Rust does not allow forcing a running thread + /// to terminate. The task will continue running until it completes or + /// encounters an error. + /// + /// Users should ensure that their blocking function periodically checks for + /// errors calling `tx.blocking_send`. An error signals that the stream has + /// been dropped / cancelled and the blocking task should exit. + /// + /// This is often used to spawn tasks that write to the sender + /// retrieved from [`Self::tx`], for examples, see the document + /// of this type. + pub fn spawn_blocking<F>(&mut self, f: F) + where + F: FnOnce() -> Result<()>, + F: Send + 'static, + { + self.inner.spawn_blocking(f) + } + + /// Same as [`Self::spawn_blocking`] but it spawns the blocking task on the provided runtime. + pub fn spawn_blocking_on<F>(&mut self, f: F, handle: &Handle) + where + F: FnOnce() -> Result<()>, + F: Send + 'static, + { + self.inner.spawn_blocking_on(f, handle) + } + + /// Runs the `partition` of the `input` ExecutionPlan on the + /// tokio thread pool and writes its outputs to this stream + /// + /// If the input partition produces an error, the error will be + /// sent to the output stream and no further results are sent. + pub(crate) fn run_input( + &mut self, + input: Arc<dyn ExecutionPlan>, + partition: usize, + context: Arc<TaskContext>, + ) { + let output = self.tx(); + let input_display = if log::log_enabled!(log::Level::Debug) { + displayable(input.as_ref()).one_line().to_string() + } else { + String::new() + }; + + self.inner.spawn(async move { + let mut stream = match input.execute(partition, context) { + Err(e) => { + // If send fails, the plan being torn down, there + // is no place to send the error and no reason to continue. + output.send(Err(e)).await.ok(); + debug!( + "Stopping execution: error executing input: {input_display}", + ); + return Ok(()); + } + Ok(stream) => stream, + }; + + // Drop the input early, as soon as we're done with it. + // Holding on to it can cause delays in cancelling the child plan when the query is + // cancelled. + drop(input); + + // Transfer batches from inner stream to the output tx + // immediately. + while let Some(item) = stream.next().await { + let is_err = item.is_err(); + + // If send fails, plan being torn down, there is no + // place to send the error and no reason to continue. + if output.send(item).await.is_err() { + debug!( + "Stopping execution: output is gone, plan cancelling: {input_display}", + ); + return Ok(()); + } + + // Stop after the first error is encountered (Don't + // drive all streams to completion) + if is_err { + debug!("Stopping execution: plan returned error: {input_display}"); + return Ok(()); + } + } + + Ok(()) + }); + } + + /// Create a stream of all [`RecordBatch`] written to `tx` + pub fn build(self) -> SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new( + self.schema, + self.inner.build(), + )) + } +} + +#[doc(hidden)] +pub struct RecordBatchReceiverStream {} + +impl RecordBatchReceiverStream { + /// Create a builder with an internal buffer of capacity batches. + pub fn builder( + schema: SchemaRef, + capacity: usize, + ) -> RecordBatchReceiverStreamBuilder { + RecordBatchReceiverStreamBuilder::new(schema, capacity) + } +} + +pin_project! { + /// Combines a [`Stream`] with a [`SchemaRef`] implementing + /// [`SendableRecordBatchStream`] for the combination + /// + /// See [`Self::new`] for an example + pub struct RecordBatchStreamAdapter<S> { + schema: SchemaRef, + + // Wrapped in Option so we can drop the inner stream as soon as it + // returns `None`, releasing any upstream pipeline resources before the + // adapter itself is dropped. + #[pin] + stream: Option<S>, + } +} + +impl<S> RecordBatchStreamAdapter<S> { + /// Creates a new [`RecordBatchStreamAdapter`] from the provided schema and stream. + /// + /// Note to create a [`SendableRecordBatchStream`] you pin the result + /// + /// # Example + /// ``` + /// # use arrow::array::record_batch; + /// # use datafusion_execution::SendableRecordBatchStream; + /// # use datafusion_physical_plan::stream::RecordBatchStreamAdapter; + /// // Create stream of Result<RecordBatch> + /// let batch = record_batch!( + /// ("a", Int32, [1, 2, 3]), + /// ("b", Float64, [Some(4.0), None, Some(5.0)]) + /// ) + /// .expect("created batch"); + /// let schema = batch.schema(); + /// let stream = futures::stream::iter(vec![Ok(batch)]); + /// // Convert the stream to a SendableRecordBatchStream + /// let adapter = RecordBatchStreamAdapter::new(schema, stream); + /// // Now you can use the adapter as a SendableRecordBatchStream + /// let batch_stream: SendableRecordBatchStream = Box::pin(adapter); + /// // ... + /// ``` + pub fn new(schema: SchemaRef, stream: S) -> Self { + Self { + schema, + stream: Some(stream), + } + } +} + +impl<S> std::fmt::Debug for RecordBatchStreamAdapter<S> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RecordBatchStreamAdapter") + .field("schema", &self.schema) + .finish() + } +} + +impl<S> Stream for RecordBatchStreamAdapter<S> +where + S: Stream<Item = Result<RecordBatch>>, +{ + type Item = Result<RecordBatch>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let mut this = self.project(); + let Some(inner) = this.stream.as_mut().as_pin_mut() else { + return Poll::Ready(None); + }; + let item = ready!(inner.poll_next(cx)); + if item.is_none() { + // Drop the inner stream in place to release its resources. + // SAFETY: the inner stream is dropped without moving it out of + // its pinned location; assigning `None` only runs the inner + // value's destructor in place, which is permitted for pinned + // values. + unsafe { + *this.stream.as_mut().get_unchecked_mut() = None; + } + } + Poll::Ready(item) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + match self.stream.as_ref() { + Some(stream) => stream.size_hint(), + None => (0, Some(0)), + } + } +} + +impl<S> RecordBatchStream for RecordBatchStreamAdapter<S> +where + S: Stream<Item = Result<RecordBatch>>, +{ + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// `EmptyRecordBatchStream` can be used to create a [`RecordBatchStream`] +/// that will produce no results +pub struct EmptyRecordBatchStream { + /// Schema wrapped by Arc + schema: SchemaRef, +} + +impl EmptyRecordBatchStream { + /// Create an empty RecordBatchStream + pub fn new(schema: SchemaRef) -> Self { + Self { schema } + } +} + +impl RecordBatchStream for EmptyRecordBatchStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for EmptyRecordBatchStream { + type Item = Result<RecordBatch>; + + fn poll_next( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + Poll::Ready(None) + } +} + +/// Stream wrapper that records `BaselineMetrics` for a particular +/// `[SendableRecordBatchStream]` (likely a partition) +pub(crate) struct ObservedStream { + inner: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, + fetch: Option<usize>, + produced: usize, +} + +impl ObservedStream { + pub fn new( + inner: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, + fetch: Option<usize>, + ) -> Self { + Self { + inner, + baseline_metrics, + fetch, + produced: 0, + } + } + + fn limit_reached( + &mut self, + poll: Poll<Option<Result<RecordBatch>>>, + ) -> Poll<Option<Result<RecordBatch>>> { + let Some(fetch) = self.fetch else { return poll }; + + if self.produced >= fetch { + self.release_inner(); + return Poll::Ready(None); + } + + if let Poll::Ready(Some(Ok(batch))) = &poll { + if self.produced + batch.num_rows() > fetch { + let batch = batch.slice(0, fetch.saturating_sub(self.produced)); + self.produced += batch.num_rows(); + if self.produced >= fetch { + self.release_inner(); + } + return Poll::Ready(Some(Ok(batch))); + }; + self.produced += batch.num_rows() + } + poll + } + + /// Replace the inner stream with an [`EmptyRecordBatchStream`], dropping + /// the original stream so its upstream pipeline can be torn down. + fn release_inner(&mut self) { + let schema = self.inner.schema(); + self.inner = Box::pin(EmptyRecordBatchStream::new(schema)); + } +} + +impl RecordBatchStream for ObservedStream { + fn schema(&self) -> SchemaRef { + self.inner.schema() + } +} + +impl Stream for ObservedStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let mut poll = self.inner.poll_next_unpin(cx); + if self.fetch.is_some() { + poll = self.limit_reached(poll); + } + self.baseline_metrics.record_poll(poll) + } +} + +pin_project! { + /// Stream wrapper that splits large [`RecordBatch`]es into smaller batches. + /// + /// This ensures upstream operators receive batches no larger than + /// `batch_size`, which can improve parallelism when data sources + /// generate very large batches. + /// + /// # Fields + /// + /// - `current_batch`: The batch currently being split, if any + /// - `offset`: Index of the next row to split from `current_batch`. + /// This tracks our position within the current batch being split. + /// + /// # Invariants + /// + /// - `offset` is always ≤ `current_batch.num_rows()` when `current_batch` is `Some` + /// - When `current_batch` is `None`, `offset` is always 0 + /// - `batch_size` is always > 0 +pub struct BatchSplitStream { + #[pin] + input: SendableRecordBatchStream, + schema: SchemaRef, + batch_size: usize, + metrics: SplitMetrics, + current_batch: Option<RecordBatch>, + offset: usize, + } +} + +impl BatchSplitStream { + /// Create a new [`BatchSplitStream`] + pub fn new( + input: SendableRecordBatchStream, + batch_size: usize, + metrics: SplitMetrics, + ) -> Self { + let schema = input.schema(); + Self { + input, + schema, + batch_size, + metrics, + current_batch: None, + offset: 0, + } + } + + /// Attempt to produce the next sliced batch from the current batch. + /// + /// Returns `Some(batch)` if a slice was produced, `None` if the current batch + /// is exhausted and we need to poll upstream for more data. + fn next_sliced_batch(&mut self) -> Option<Result<RecordBatch>> { + let batch = self.current_batch.take()?; + + // Assert slice boundary safety - offset should never exceed batch size + debug_assert!( + self.offset <= batch.num_rows(), + "Offset {} exceeds batch size {}", + self.offset, + batch.num_rows() + ); + + let remaining = batch.num_rows() - self.offset; + let to_take = remaining.min(self.batch_size); + let out = batch.slice(self.offset, to_take); + + self.metrics.batches_split.add(1); + self.offset += to_take; + if self.offset < batch.num_rows() { + // More data remains in this batch, store it back + self.current_batch = Some(batch); + } else { + // Batch is exhausted, reset offset + // Note: current_batch is already None since we took it at the start + self.offset = 0; + } + Some(Ok(out)) + } + + /// Poll the upstream input for the next batch. + /// + /// Returns the appropriate `Poll` result based on upstream state. + /// Small batches are passed through directly, large batches are stored + /// for slicing and return the first slice immediately. + fn poll_upstream( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + match ready!(self.input.as_mut().poll_next(cx)) { + Some(Ok(batch)) => { + if batch.num_rows() <= self.batch_size { + // Small batch, pass through directly + Poll::Ready(Some(Ok(batch))) + } else { + // Large batch, store for slicing and return first slice + self.current_batch = Some(batch); + // Immediately produce the first slice + match self.next_sliced_batch() { + Some(result) => Poll::Ready(Some(result)), + None => Poll::Ready(None), // Should not happen + } + } + } + Some(Err(e)) => Poll::Ready(Some(Err(e))), + None => { + // Release the input pipeline's resources. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + Poll::Ready(None) + } + } + } +} + +impl Stream for BatchSplitStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + // First, try to produce a slice from the current batch + if let Some(result) = self.next_sliced_batch() { + return Poll::Ready(Some(result)); + } + + // No current batch or current batch exhausted, poll upstream + self.poll_upstream(cx) + } +} + +impl RecordBatchStream for BatchSplitStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// A stream that holds a memory reservation for its lifetime, +/// shrinking the reservation as batches are consumed. +/// The original reservation must have its batch sizes calculated using [`get_record_batch_memory_size`] +/// On error, the reservation is *NOT* freed, until the stream is dropped. +pub(crate) struct ReservationStream { + schema: SchemaRef, + inner: SendableRecordBatchStream, + reservation: MemoryReservation, +} + +impl ReservationStream { + pub(crate) fn new( + schema: SchemaRef, + inner: SendableRecordBatchStream, + reservation: MemoryReservation, + ) -> Self { + Self { + schema, + inner, + reservation, + } + } +} + +impl Stream for ReservationStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let res = self.inner.poll_next_unpin(cx); + + match res { + Poll::Ready(res) => { + match res { + Some(Ok(batch)) => { + self.reservation + .shrink(get_record_batch_memory_size(&batch)); + Poll::Ready(Some(Ok(batch))) + } + Some(Err(err)) => Poll::Ready(Some(Err(err))), + None => { + // Stream is done so free the reservation completely + self.reservation.free(); + // Release the input pipeline's resources. + let inner_schema = self.inner.schema(); + self.inner = Box::pin(EmptyRecordBatchStream::new(inner_schema)); + Poll::Ready(None) + } + } + } + Poll::Pending => Poll::Pending, + } + } + + fn size_hint(&self) -> (usize, Option<usize>) { + self.inner.size_hint() + } +} + +impl RecordBatchStream for ReservationStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::test::exec::{ + BlockingExec, MockExec, PanicExec, assert_strong_count_converges_to_zero, + }; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::exec_err; + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])) + } + + #[tokio::test] + #[should_panic(expected = "PanickingStream did panic")] + async fn record_batch_receiver_stream_propagates_panics() { + let schema = schema(); + + let num_partitions = 10; + let input = PanicExec::new(Arc::clone(&schema), num_partitions); + consume(input, 10).await + } + + #[tokio::test] + #[should_panic(expected = "PanickingStream did panic: 1")] + async fn record_batch_receiver_stream_propagates_panics_early_shutdown() { + let schema = schema(); + + // Make 2 partitions, second partition panics before the first + let num_partitions = 2; + let input = PanicExec::new(Arc::clone(&schema), num_partitions) + .with_partition_panic(0, 10) + .with_partition_panic(1, 3); // partition 1 should panic first (after 3 ) + + // Ensure that the panic results in an early shutdown (that + // everything stops after the first panic). + + // Since the stream reads every other batch: (0,1,0,1,0,panic) + // so should not exceed 5 batches prior to the panic + let max_batches = 5; + consume(input, max_batches).await + } + + #[tokio::test] + async fn record_batch_receiver_stream_drop_cancel() { + let task_ctx = Arc::new(TaskContext::default()); + let schema = schema(); + + // Make an input that never proceeds + let input = BlockingExec::new(Arc::clone(&schema), 1); + let refs = input.refs(); + + // Configure a RecordBatchReceiverStream to consume the input + let mut builder = RecordBatchReceiverStream::builder(schema, 2); + builder.run_input(Arc::new(input), 0, Arc::clone(&task_ctx)); + let stream = builder.build(); + + // Input should still be present + assert!(std::sync::Weak::strong_count(&refs) > 0); + + // Drop the stream, ensure the refs go to zero + drop(stream); + assert_strong_count_converges_to_zero(refs).await; + } + + #[tokio::test] + /// Ensure that if an error is received in one stream, the + /// `RecordBatchReceiverStream` stops early and does not drive + /// other streams to completion. + async fn record_batch_receiver_stream_error_does_not_drive_completion() { + let task_ctx = Arc::new(TaskContext::default()); + let schema = schema(); + + // make an input that will error twice + let error_stream = MockExec::new( + vec![exec_err!("Test1"), exec_err!("Test2")], + Arc::clone(&schema), + ) + .with_use_task(false); + + let mut builder = RecordBatchReceiverStream::builder(schema, 2); + builder.run_input(Arc::new(error_stream), 0, Arc::clone(&task_ctx)); + let mut stream = builder.build(); + + // Get the first result, which should be an error + let first_batch = stream.next().await.unwrap(); + let first_err = first_batch.unwrap_err(); + assert_eq!(first_err.strip_backtrace(), "Execution error: Test1"); + + // There should be no more batches produced (should not get the second error) + assert!(stream.next().await.is_none()); + } + + #[tokio::test] + async fn batch_split_stream_basic_functionality() { + use arrow::array::{Int32Array, RecordBatch}; + use futures::stream::{self, StreamExt}; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create a large batch that should be split + let large_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from((0..2000).collect::<Vec<_>>()))], + ) + .unwrap(); + + // Create a stream with the large batch + let input_stream = stream::iter(vec![Ok(large_batch)]); + let adapter = RecordBatchStreamAdapter::new(Arc::clone(&schema), input_stream); + let batch_stream = Box::pin(adapter) as SendableRecordBatchStream; + + // Create a BatchSplitStream with batch_size = 500 + let metrics = ExecutionPlanMetricsSet::new(); + let split_metrics = SplitMetrics::new(&metrics, 0); + let mut split_stream = BatchSplitStream::new(batch_stream, 500, split_metrics); + + let mut total_rows = 0; + let mut batch_count = 0; + + while let Some(result) = split_stream.next().await { + let batch = result.unwrap(); + assert!(batch.num_rows() <= 500, "Batch size should not exceed 500"); + total_rows += batch.num_rows(); + batch_count += 1; + } + + assert_eq!(total_rows, 2000, "All rows should be preserved"); + assert_eq!(batch_count, 4, "Should have 4 batches of 500 rows each"); + } + + /// Consumes all the input's partitions into a + /// RecordBatchReceiverStream and runs it to completion + /// + /// panic's if more than max_batches is seen, + async fn consume(input: PanicExec, max_batches: usize) { + let task_ctx = Arc::new(TaskContext::default()); + + let input = Arc::new(input); + let num_partitions = input.properties().output_partitioning().partition_count(); + + // Configure a RecordBatchReceiverStream to consume all the input partitions + let mut builder = + RecordBatchReceiverStream::builder(input.schema(), num_partitions); + for partition in 0..num_partitions { + builder.run_input( + Arc::clone(&input) as Arc<dyn ExecutionPlan>, + partition, + Arc::clone(&task_ctx), + ); + } + let mut stream = builder.build(); + + // Drain the stream until it is complete, panic'ing on error + let mut num_batches = 0; + while let Some(next) = stream.next().await { + next.unwrap(); + num_batches += 1; + assert!( + num_batches < max_batches, + "Got the limit of {num_batches} batches before seeing panic" + ); + } + } + + #[test] + fn record_batch_receiver_stream_builder_spawn_on_runtime() { + let tokio_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + + let mut builder = + RecordBatchReceiverStreamBuilder::new(Arc::new(Schema::empty()), 10); + + let tx1 = builder.tx(); + builder.spawn_on( + async move { + tx1.send(Ok(RecordBatch::new_empty(Arc::new(Schema::empty())))) + .await + .unwrap(); + + Ok(()) + }, + tokio_runtime.handle(), + ); + + let tx2 = builder.tx(); + builder.spawn_blocking_on( + move || { + tx2.blocking_send(Ok(RecordBatch::new_empty(Arc::new(Schema::empty())))) + .unwrap(); + + Ok(()) + }, + tokio_runtime.handle(), + ); + + let mut stream = builder.build(); + + let mut number_of_batches = 0; + + loop { + let poll = stream.poll_next_unpin(&mut Context::from_waker( + futures::task::noop_waker_ref(), + )); + + match poll { + Poll::Ready(None) => { + break; + } + Poll::Ready(Some(Ok(batch))) => { + number_of_batches += 1; + assert_eq!(batch.num_rows(), 0); + } + Poll::Ready(Some(Err(e))) => panic!("Unexpected error: {e}"), + Poll::Pending => { + continue; + } + } + } + + assert_eq!( + number_of_batches, 2, + "Should have received exactly two empty batches" + ); + } + + #[tokio::test] + async fn test_reservation_stream_shrinks_on_poll() { + use arrow::array::Int32Array; + use datafusion_execution::memory_pool::MemoryConsumer; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(10 * 1024 * 1024, 1.0) + .build_arc() + .unwrap(); + + let reservation = MemoryConsumer::new("test").register(&runtime.memory_pool); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + // Create batches + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + ) + .unwrap(); + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![6, 7, 8, 9, 10]))], + ) + .unwrap(); + + let batch1_size = get_record_batch_memory_size(&batch1); + let batch2_size = get_record_batch_memory_size(&batch2); + + // Reserve memory upfront + reservation.try_grow(batch1_size + batch2_size).unwrap(); + let initial_reserved = runtime.memory_pool.reserved(); + assert_eq!(initial_reserved, batch1_size + batch2_size); + + // Create stream with batches + let stream = futures::stream::iter(vec![Ok(batch1), Ok(batch2)]); + let inner = Box::pin(RecordBatchStreamAdapter::new(Arc::clone(&schema), stream)) + as SendableRecordBatchStream; + + let mut res_stream = + ReservationStream::new(Arc::clone(&schema), inner, reservation); + + // Poll first batch + let result1 = res_stream.next().await; + assert!(result1.is_some()); + + // Memory should be reduced by batch1_size + let after_first = runtime.memory_pool.reserved(); + assert_eq!(after_first, batch2_size); + + // Poll second batch + let result2 = res_stream.next().await; + assert!(result2.is_some()); + + // Memory should be reduced by batch2_size + let after_second = runtime.memory_pool.reserved(); + assert_eq!(after_second, 0); + + // Poll None (end of stream) + let result3 = res_stream.next().await; + assert!(result3.is_none()); + + // Memory should still be 0 + assert_eq!(runtime.memory_pool.reserved(), 0); + } + + #[tokio::test] + async fn test_reservation_stream_error_handling() { + use datafusion_execution::memory_pool::MemoryConsumer; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(10 * 1024 * 1024, 1.0) + .build_arc() + .unwrap(); + + let reservation = MemoryConsumer::new("test").register(&runtime.memory_pool); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + reservation.try_grow(1000).unwrap(); + let initial = runtime.memory_pool.reserved(); + assert_eq!(initial, 1000); + + // Create a stream that errors + let stream = futures::stream::iter(vec![exec_err!("Test error")]); + let inner = Box::pin(RecordBatchStreamAdapter::new(Arc::clone(&schema), stream)) + as SendableRecordBatchStream; + + let mut res_stream = + ReservationStream::new(Arc::clone(&schema), inner, reservation); + + // Get the error + let result = res_stream.next().await; + assert!(result.is_some()); + assert!(result.unwrap().is_err()); + + // Verify reservation is NOT automatically freed on error + // The reservation is only freed when poll_next returns Poll::Ready(None) + // After an error, the stream may continue to hold the reservation + // until it's explicitly dropped or polled to None + let after_error = runtime.memory_pool.reserved(); + assert_eq!( + after_error, 1000, + "Reservation should still be held after error" + ); + + // Drop the stream to free the reservation + drop(res_stream); + + // Now memory should be freed + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "Memory should be freed when stream is dropped" + ); + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/streaming.rs b/native/vendor/datafusion-physical-plan/src/streaming.rs new file mode 100644 index 0000000..7b0058e --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/streaming.rs
@@ -0,0 +1,486 @@ +// 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. + +//! Generic plans for deferred execution: [`StreamingTableExec`] and [`PartitionStream`] + +use std::fmt::Debug; +use std::sync::Arc; + +use super::{DisplayAs, DisplayFormatType, PlanProperties}; +use crate::coop::make_cooperative; +use crate::display::{ProjectSchemaDisplay, display_orderings}; +use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; +use crate::limit::LimitStream; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::projection::{ + ProjectionExec, all_alias_free_columns, new_projections_for_columns, update_ordering, +}; +use crate::stream::RecordBatchStreamAdapter; +use crate::{ + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; + +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, internal_err, plan_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; + +use async_trait::async_trait; +use futures::stream::StreamExt; +use log::debug; + +/// A partition that can be converted into a [`SendableRecordBatchStream`] +/// +/// Combined with [`StreamingTableExec`], you can use this trait to implement +/// [`ExecutionPlan`] for a custom source with less boiler plate than +/// implementing `ExecutionPlan` directly for many use cases. +pub trait PartitionStream: Debug + Send + Sync { + /// Returns the schema of this partition + fn schema(&self) -> &SchemaRef; + + /// Returns a stream yielding this partitions values + fn execute(&self, ctx: Arc<TaskContext>) -> SendableRecordBatchStream; +} + +/// An [`ExecutionPlan`] for one or more [`PartitionStream`]s. +/// +/// If your source can be represented as one or more [`PartitionStream`]s, you can +/// use this struct to implement [`ExecutionPlan`]. +#[derive(Clone)] +pub struct StreamingTableExec { + partitions: Vec<Arc<dyn PartitionStream>>, + projection: Option<Arc<[usize]>>, + projected_schema: SchemaRef, + projected_output_ordering: Vec<LexOrdering>, + infinite: bool, + limit: Option<usize>, + cache: Arc<PlanProperties>, + metrics: ExecutionPlanMetricsSet, +} + +impl StreamingTableExec { + /// Try to create a new [`StreamingTableExec`] returning an error if the schema is incorrect + pub fn try_new( + schema: SchemaRef, + partitions: Vec<Arc<dyn PartitionStream>>, + projection: Option<&Vec<usize>>, + projected_output_ordering: impl IntoIterator<Item = LexOrdering>, + infinite: bool, + limit: Option<usize>, + ) -> Result<Self> { + for x in partitions.iter() { + let partition_schema = x.schema(); + if !schema.eq(partition_schema) { + debug!( + "Target schema does not match with partition schema. \ + Target_schema: {schema:?}. Partition Schema: {partition_schema:?}" + ); + return plan_err!("Mismatch between schema and batches"); + } + } + + let projected_schema = match projection { + Some(p) => Arc::new(schema.project(p)?), + None => schema, + }; + let projected_output_ordering = + projected_output_ordering.into_iter().collect::<Vec<_>>(); + let cache = Self::compute_properties( + Arc::clone(&projected_schema), + projected_output_ordering.clone(), + Partitioning::UnknownPartitioning(partitions.len()), + infinite, + ); + Ok(Self { + partitions, + projected_schema, + projection: projection.cloned().map(Into::into), + projected_output_ordering, + infinite, + limit, + cache: Arc::new(cache), + metrics: ExecutionPlanMetricsSet::new(), + }) + } + + /// Declares the output partitioning of this stream. + /// + /// `output_partitioning` must describe this plan's current output and have + /// the same number of partitions as the stream. + pub fn with_output_partitioning( + mut self, + output_partitioning: Partitioning, + ) -> Result<Self> { + if output_partitioning.partition_count() != self.partitions.len() { + return plan_err!( + "Output partitioning has {} partitions but stream has {} partitions", + output_partitioning.partition_count(), + self.partitions.len() + ); + } + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + Ok(self) + } + + pub fn partitions(&self) -> &Vec<Arc<dyn PartitionStream>> { + &self.partitions + } + + pub fn partition_schema(&self) -> &SchemaRef { + self.partitions[0].schema() + } + + pub fn projection(&self) -> &Option<Arc<[usize]>> { + &self.projection + } + + pub fn projected_schema(&self) -> &Schema { + &self.projected_schema + } + + pub fn projected_output_ordering(&self) -> impl IntoIterator<Item = LexOrdering> { + self.projected_output_ordering.clone() + } + + pub fn is_infinite(&self) -> bool { + self.infinite + } + + pub fn limit(&self) -> Option<usize> { + self.limit + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + orderings: Vec<LexOrdering>, + output_partitioning: Partitioning, + infinite: bool, + ) -> PlanProperties { + // Calculate equivalence properties: + let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings); + + let boundedness = if infinite { + Boundedness::Unbounded { + requires_infinite_memory: false, + } + } else { + Boundedness::Bounded + }; + PlanProperties::new( + eq_properties, + output_partitioning, + EmissionType::Incremental, + boundedness, + ) + .with_scheduling_type(SchedulingType::Cooperative) + } +} + +impl Debug for StreamingTableExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LazyMemTableExec").finish_non_exhaustive() + } +} + +impl DisplayAs for StreamingTableExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "StreamingTableExec: partition_sizes={:?}", + self.partitions.len(), + )?; + if !self.projected_schema.fields().is_empty() { + write!( + f, + ", projection={}", + ProjectSchemaDisplay(&self.projected_schema) + )?; + } + if self.infinite { + write!(f, ", infinite_source=true")?; + } + if let Some(fetch) = self.limit { + write!(f, ", fetch={fetch}")?; + } + if !matches!( + self.cache.output_partitioning(), + Partitioning::UnknownPartitioning(_) + ) { + write!( + f, + ", output_partitioning={}", + self.cache.output_partitioning() + )?; + } + + display_orderings(f, &self.projected_output_ordering)?; + + Ok(()) + } + DisplayFormatType::TreeRender => { + if self.infinite { + writeln!(f, "infinite={}", self.infinite)?; + } + if let Some(limit) = self.limit { + write!(f, "limit={limit}")?; + } else { + write!(f, "limit=None")?; + } + + Ok(()) + } + } + } +} + +#[async_trait] +impl ExecutionPlan for StreamingTableExec { + fn name(&self) -> &'static str { + "StreamingTableExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn fetch(&self) -> Option<usize> { + self.limit + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + if children.is_empty() { + Ok(self) + } else { + internal_err!("Children cannot be replaced in {self:?}") + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + ctx: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let stream = self.partitions[partition].execute(Arc::clone(&ctx)); + let projected_stream = match self.projection.clone() { + Some(projection) => Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.projected_schema), + stream.map(move |x| { + x.and_then(|b| b.project(projection.as_ref()).map_err(Into::into)) + }), + )), + None => stream, + }; + let stream = make_cooperative(projected_stream); + + Ok(match self.limit { + None => stream, + Some(fetch) => { + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + Box::pin(LimitStream::new(stream, 0, Some(fetch), baseline_metrics)) + } + }) + } + + /// Tries to embed `projection` to its input (`streaming table`). + /// If possible, returns [`StreamingTableExec`] as the top plan. Otherwise, + /// returns `None`. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + if !all_alias_free_columns(projection.expr()) { + return Ok(None); + } + + let streaming_table_projections = + self.projection().as_ref().map(|i| i.as_ref().to_vec()); + let new_projections = new_projections_for_columns( + projection.expr(), + &streaming_table_projections + .unwrap_or_else(|| (0..self.schema().fields().len()).collect()), + ); + + let mut lex_orderings = vec![]; + for ordering in self.projected_output_ordering().into_iter() { + let Some(ordering) = update_ordering(ordering, projection.expr())? else { + return Ok(None); + }; + lex_orderings.push(ordering); + } + let projection_mapping = ProjectionMapping::try_new( + projection + .expr() + .iter() + .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())), + &self.schema(), + )?; + let output_partitioning = self + .cache + .output_partitioning() + .project(&projection_mapping, self.cache.equivalence_properties()); + + StreamingTableExec::try_new( + Arc::clone(self.partition_schema()), + self.partitions().clone(), + Some(new_projections.as_ref()), + lex_orderings, + self.is_infinite(), + self.limit(), + ) + .and_then(|exec| exec.with_output_partitioning(output_partitioning)) + .map(|e| Some(Arc::new(e) as _)) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> { + Some(Arc::new(StreamingTableExec { + partitions: self.partitions.clone(), + projection: self.projection.clone(), + projected_schema: Arc::clone(&self.projected_schema), + projected_output_ordering: self.projected_output_ordering.clone(), + infinite: self.infinite, + limit, + cache: Arc::clone(&self.cache), + metrics: self.metrics.clone(), + })) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::collect_partitioned; + use crate::streaming::PartitionStream; + use crate::test::{TestPartitionStream, make_partition}; + use arrow::record_batch::RecordBatch; + + #[tokio::test] + async fn test_no_limit() { + let exec = TestBuilder::new() + // Make 2 batches, each with 100 rows + .with_batches(vec![make_partition(100), make_partition(100)]) + .build(); + + let counts = collect_num_rows(Arc::new(exec)).await; + assert_eq!(counts, vec![200]); + } + + #[tokio::test] + async fn test_limit() { + let exec = TestBuilder::new() + // Make 2 batches, each with 100 rows + .with_batches(vec![make_partition(100), make_partition(100)]) + // Limit to only the first 75 rows back + .with_limit(Some(75)) + .build(); + + let counts = collect_num_rows(Arc::new(exec)).await; + assert_eq!(counts, vec![75]); + } + + /// Runs the provided execution plan and returns a vector of the number of + /// rows in each partition + async fn collect_num_rows(exec: Arc<dyn ExecutionPlan>) -> Vec<usize> { + let ctx = Arc::new(TaskContext::default()); + let partition_batches = collect_partitioned(exec, ctx).await.unwrap(); + partition_batches + .into_iter() + .map(|batches| batches.iter().map(|b| b.num_rows()).sum::<usize>()) + .collect() + } + + #[derive(Default)] + struct TestBuilder { + schema: Option<SchemaRef>, + partitions: Vec<Arc<dyn PartitionStream>>, + projection: Option<Vec<usize>>, + projected_output_ordering: Vec<LexOrdering>, + infinite: bool, + limit: Option<usize>, + } + + impl TestBuilder { + fn new() -> Self { + Self::default() + } + + /// Set the batches for the stream + fn with_batches(mut self, batches: Vec<RecordBatch>) -> Self { + let stream = TestPartitionStream::new_with_batches(batches); + self.schema = Some(Arc::clone(stream.schema())); + self.partitions = vec![Arc::new(stream)]; + self + } + + /// Set the limit for the stream + fn with_limit(mut self, limit: Option<usize>) -> Self { + self.limit = limit; + self + } + + fn build(self) -> StreamingTableExec { + StreamingTableExec::try_new( + self.schema.unwrap(), + self.partitions, + self.projection.as_ref(), + self.projected_output_ordering, + self.infinite, + self.limit, + ) + .unwrap() + } + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/test.rs b/native/vendor/datafusion-physical-plan/src/test.rs new file mode 100644 index 0000000..b38a46d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/test.rs
@@ -0,0 +1,570 @@ +// 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. + +//! Utilities for testing datafusion-physical-plan + +use std::collections::HashMap; +use std::fmt; +use std::fmt::{Debug, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::Context; + +use crate::common; +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::memory::MemoryStream; +use crate::metrics::MetricsSet; +use crate::statistics::StatisticsArgs; +use crate::stream::RecordBatchStreamAdapter; +use crate::streaming::PartitionStream; +use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; +use crate::{DisplayAs, DisplayFormatType, PlanProperties}; + +use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + Result, Statistics, assert_or_internal_err, config::ConfigOptions, project_schema, +}; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_expr::equivalence::{ + OrderingEquivalenceClass, ProjectionMapping, +}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, PhysicalExpr, +}; + +use futures::{Future, FutureExt}; + +pub mod exec; + +/// `TestMemoryExec` is a mock equivalent to [`MemorySourceConfig`] with [`ExecutionPlan`] implemented for testing. +/// i.e. It has some but not all the functionality of [`MemorySourceConfig`]. +/// This implements an in-memory DataSource rather than explicitly implementing a trait. +/// It is implemented in this manner to keep relevant unit tests in place +/// while avoiding circular dependencies between `datafusion-physical-plan` and `datafusion-datasource`. +/// +/// [`MemorySourceConfig`]: https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs +#[derive(Clone, Debug)] +pub struct TestMemoryExec { + /// The partitions to query + partitions: Vec<Vec<RecordBatch>>, + /// Schema representing the data before projection + schema: SchemaRef, + /// Schema representing the data after the optional projection is applied + projected_schema: SchemaRef, + /// Optional projection + projection: Option<Vec<usize>>, + /// Sort information: one or more equivalent orderings + sort_information: Vec<LexOrdering>, + /// if partition sizes should be displayed + show_sizes: bool, + /// The maximum number of records to read from this plan. If `None`, + /// all records after filtering are returned. + fetch: Option<usize>, + cache: Arc<PlanProperties>, +} + +impl DisplayAs for TestMemoryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "DataSourceExec: ")?; + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let partition_sizes: Vec<_> = + self.partitions.iter().map(|b| b.len()).collect(); + + let output_ordering = self + .sort_information + .first() + .map(|output_ordering| format!(", output_ordering={output_ordering}")) + .unwrap_or_default(); + + let eq_properties = self.eq_properties(); + let constraints = eq_properties.constraints(); + let constraints = if constraints.is_empty() { + String::new() + } else { + format!(", {constraints}") + }; + + let limit = self + .fetch + .map_or(String::new(), |limit| format!(", fetch={limit}")); + if self.show_sizes { + write!( + f, + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) + } else { + write!( + f, + "partitions={}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) + } + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for TestMemoryExec { + fn name(&self) -> &'static str { + "DataSourceExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn repartitioned( + &self, + _target_partitions: usize, + _config: &ConfigOptions, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + unimplemented!() + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + self.open(partition, context) + } + + fn metrics(&self) -> Option<MetricsSet> { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if args.partition().is_some() { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } else { + Ok(Arc::new(self.statistics_inner()?)) + } + } + + fn fetch(&self) -> Option<usize> { + self.fetch + } +} + +impl TestMemoryExec { + fn open( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin( + MemoryStream::try_new( + self.partitions[partition].clone(), + Arc::clone(&self.projected_schema), + self.projection.clone(), + )? + .with_fetch(self.fetch), + )) + } + + fn compute_properties(&self) -> PlanProperties { + PlanProperties::new( + self.eq_properties(), + self.output_partitioning(), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(self.partitions.len()) + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new_with_orderings( + Arc::clone(&self.projected_schema), + self.sort_information.clone(), + ) + } + + fn statistics_inner(&self) -> Result<Statistics> { + Ok(common::compute_record_batch_statistics( + &self.partitions, + &self.schema, + self.projection.clone(), + )) + } + + pub fn try_new( + partitions: &[Vec<RecordBatch>], + schema: SchemaRef, + projection: Option<Vec<usize>>, + ) -> Result<Self> { + let projected_schema = project_schema(&schema, projection.as_ref())?; + Ok(Self { + partitions: partitions.to_vec(), + schema, + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings( + Arc::clone(&projected_schema), + Vec::<LexOrdering>::new(), + ), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )), + projected_schema, + projection, + sort_information: vec![], + show_sizes: true, + fetch: None, + }) + } + + /// Create a new `DataSourceExec` Equivalent plan for reading in-memory record batches + /// The provided `schema` should not have the projection applied. + pub fn try_new_exec( + partitions: &[Vec<RecordBatch>], + schema: SchemaRef, + projection: Option<Vec<usize>>, + ) -> Result<Arc<TestMemoryExec>> { + let mut source = Self::try_new(partitions, schema, projection)?; + let cache = source.compute_properties(); + source.cache = Arc::new(cache); + Ok(Arc::new(source)) + } + + // Equivalent of `DataSourceExec::new` + pub fn update_cache(source: &Arc<TestMemoryExec>) -> TestMemoryExec { + let cache = source.compute_properties(); + let mut source = (**source).clone(); + source.cache = Arc::new(cache); + source + } + + /// Set the limit of the files + pub fn with_limit(mut self, limit: Option<usize>) -> Self { + self.fetch = limit; + self + } + + /// Ref to partitions + pub fn partitions(&self) -> &[Vec<RecordBatch>] { + &self.partitions + } + + /// Ref to projection + pub fn projection(&self) -> &Option<Vec<usize>> { + &self.projection + } + + /// Ref to sort information + pub fn sort_information(&self) -> &[LexOrdering] { + &self.sort_information + } + + /// refer to `try_with_sort_information` at MemorySourceConfig for more information. + /// <https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs> + pub fn try_with_sort_information( + mut self, + mut sort_information: Vec<LexOrdering>, + ) -> Result<Self> { + // All sort expressions must refer to the original schema + let fields = self.schema.fields(); + let ambiguous_column = sort_information + .iter() + .flat_map(|ordering| ordering.clone()) + .flat_map(|expr| collect_columns(&expr.expr)) + .find(|col| { + fields + .get(col.index()) + .map(|field| field.name() != col.name()) + .unwrap_or(true) + }); + assert_or_internal_err!( + ambiguous_column.is_none(), + "Column {:?} is not found in the original schema of the TestMemoryExec", + ambiguous_column.as_ref().unwrap() + ); + + // If there is a projection on the source, we also need to project orderings + if let Some(projection) = &self.projection { + let base_schema = self.original_schema(); + let proj_exprs = projection.iter().map(|idx| { + let name = base_schema.field(*idx).name(); + (Arc::new(Column::new(name, *idx)) as _, name.to_string()) + }); + let projection_mapping = + ProjectionMapping::try_new(proj_exprs, &base_schema)?; + let base_eqp = EquivalenceProperties::new_with_orderings( + Arc::clone(&base_schema), + sort_information, + ); + let proj_eqp = + base_eqp.project(&projection_mapping, Arc::clone(&self.projected_schema)); + let oeq_class: OrderingEquivalenceClass = proj_eqp.into(); + sort_information = oeq_class.into(); + } + + self.sort_information = sort_information; + self.cache = Arc::new(self.compute_properties()); + Ok(self) + } + + /// Arc clone of ref to original schema + pub fn original_schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Asserts that given future is pending. +pub fn assert_is_pending<'a, T>(fut: &mut Pin<Box<dyn Future<Output = T> + Send + 'a>>) { + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + let poll = fut.poll_unpin(&mut cx); + + assert!(poll.is_pending()); +} + +/// Get the schema for the aggregate_test_* csv files +pub fn aggr_test_schema() -> SchemaRef { + let mut f1 = Field::new("c1", DataType::Utf8, false); + f1.set_metadata(HashMap::from_iter(vec![("testing".into(), "test".into())])); + let schema = Schema::new(vec![ + f1, + Field::new("c2", DataType::UInt32, false), + Field::new("c3", DataType::Int8, false), + Field::new("c4", DataType::Int16, false), + Field::new("c5", DataType::Int32, false), + Field::new("c6", DataType::Int64, false), + Field::new("c7", DataType::UInt8, false), + Field::new("c8", DataType::UInt16, false), + Field::new("c9", DataType::UInt32, false), + Field::new("c10", DataType::UInt64, false), + Field::new("c11", DataType::Float32, false), + Field::new("c12", DataType::Float64, false), + Field::new("c13", DataType::Utf8, false), + ]); + + Arc::new(schema) +} + +/// Returns record batch with 3 columns of i32 in memory +pub fn build_table_i32( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), +) -> RecordBatch { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Int32, false), + Field::new(b.0, DataType::Int32, false), + Field::new(c.0, DataType::Int32, false), + ]); + + RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int32Array::from(a.1.clone())), + Arc::new(Int32Array::from(b.1.clone())), + Arc::new(Int32Array::from(c.1.clone())), + ], + ) + .unwrap() +} + +/// Returns record batch with 2 columns of i32 in memory +pub fn build_table_i32_two_cols( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), +) -> RecordBatch { + let schema = Schema::new(vec![ + Field::new(a.0, DataType::Int32, false), + Field::new(b.0, DataType::Int32, false), + ]); + + RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int32Array::from(a.1.clone())), + Arc::new(Int32Array::from(b.1.clone())), + ], + ) + .unwrap() +} + +/// Returns memory table scan wrapped around record batch with 3 columns of i32 +pub fn build_table_scan_i32( + a: (&str, &Vec<i32>), + b: (&str, &Vec<i32>), + c: (&str, &Vec<i32>), +) -> Arc<dyn ExecutionPlan> { + let batch = build_table_i32(a, b, c); + let schema = batch.schema(); + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +/// Return a RecordBatch with a single Int32 array with values (0..sz) in a field named "i" +pub fn make_partition(sz: i32) -> RecordBatch { + let seq_start = 0; + let seq_end = sz; + let values = (seq_start..seq_end).collect::<Vec<_>>(); + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let arr = Arc::new(Int32Array::from(values)); + let arr = arr as ArrayRef; + + RecordBatch::try_new(schema, vec![arr]).unwrap() +} + +pub fn make_partition_utf8(sz: i32) -> RecordBatch { + let seq_start = 0; + let seq_end = sz; + let values = (seq_start..seq_end) + .map(|i| format!("test_long_string_that_is_roughly_42_bytes_{i}")) + .collect::<Vec<_>>(); + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Utf8, true)])); + let mut string_array = arrow::array::StringArray::from(values); + string_array.shrink_to_fit(); + let arr = Arc::new(string_array); + let arr = arr as ArrayRef; + + RecordBatch::try_new(schema, vec![arr]).unwrap() +} + +/// Returns a `DataSourceExec` that scans `partitions` of 100 batches each +pub fn scan_partitioned(partitions: usize) -> Arc<dyn ExecutionPlan> { + Arc::new(mem_exec(partitions)) +} + +pub fn scan_partitioned_utf8(partitions: usize) -> Arc<dyn ExecutionPlan> { + Arc::new(mem_exec_utf8(partitions)) +} + +/// Returns a `DataSourceExec` that scans `partitions` of 100 batches each +pub fn mem_exec(partitions: usize) -> TestMemoryExec { + let data: Vec<Vec<_>> = (0..partitions).map(|_| vec![make_partition(100)]).collect(); + + let schema = data[0][0].schema(); + let projection = None; + + TestMemoryExec::try_new(&data, schema, projection).unwrap() +} + +pub fn mem_exec_utf8(partitions: usize) -> TestMemoryExec { + let data: Vec<Vec<_>> = (0..partitions) + .map(|_| vec![make_partition_utf8(100)]) + .collect(); + + let schema = data[0][0].schema(); + let projection = None; + + TestMemoryExec::try_new(&data, schema, projection).unwrap() +} + +// Construct a stream partition for test purposes +#[derive(Debug)] +pub struct TestPartitionStream { + pub schema: SchemaRef, + pub batches: Vec<RecordBatch>, +} + +impl TestPartitionStream { + /// Create a new stream partition with the provided batches + pub fn new_with_batches(batches: Vec<RecordBatch>) -> Self { + let schema = batches[0].schema(); + Self { schema, batches } + } +} +impl PartitionStream for TestPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.schema + } + fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream { + let stream = futures::stream::iter(self.batches.clone().into_iter().map(Ok)); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + )) + } +} + +#[cfg(test)] +macro_rules! assert_join_metrics { + ($metrics:expr, $expected_rows:expr) => { + assert_eq!($metrics.output_rows().unwrap(), $expected_rows); + + let elapsed_compute = $metrics + .elapsed_compute() + .expect("did not find elapsed_compute metric"); + let join_time = $metrics + .sum_by_name("join_time") + .expect("did not find join_time metric") + .as_usize(); + let build_time = $metrics + .sum_by_name("build_time") + .expect("did not find build_time metric") + .as_usize(); + // ensure join_time and build_time are considered in elapsed_compute + assert!( + join_time + build_time <= elapsed_compute, + "join_time ({}) + build_time ({}) = {} was <= elapsed_compute = {}", + join_time, + build_time, + join_time + build_time, + elapsed_compute + ); + }; +} +#[cfg(test)] +pub(crate) use assert_join_metrics;
diff --git a/native/vendor/datafusion-physical-plan/src/test/exec.rs b/native/vendor/datafusion-physical-plan/src/test/exec.rs new file mode 100644 index 0000000..1e2005e --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/test/exec.rs
@@ -0,0 +1,1087 @@ +// 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. + +//! Simple iterator over batches for use in testing + +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, Statistics, common, + execution_plan::Boundedness, statistics::StatisticsArgs, +}; +use crate::{ + execution_plan::EmissionType, + stream::{RecordBatchReceiverStream, RecordBatchStreamAdapter}, +}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::{ + pin::Pin, + sync::{Arc, Weak}, + task::{Context, Poll}, +}; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; + +use futures::Stream; +use tokio::sync::Barrier; + +/// Index into the data that has been returned so far +#[derive(Debug, Default, Clone)] +pub struct BatchIndex { + inner: Arc<std::sync::Mutex<usize>>, +} + +impl BatchIndex { + /// Return the current index + pub fn value(&self) -> usize { + let inner = self.inner.lock().unwrap(); + *inner + } + + // increment the current index by one + pub fn incr(&self) { + let mut inner = self.inner.lock().unwrap(); + *inner += 1; + } +} + +/// Iterator over batches +#[derive(Debug, Default)] +pub struct TestStream { + /// Vector of record batches + data: Vec<RecordBatch>, + /// Index into the data that has been returned so far + index: BatchIndex, +} + +impl TestStream { + /// Create an iterator for a vector of record batches. Assumes at + /// least one entry in data (for the schema) + pub fn new(data: Vec<RecordBatch>) -> Self { + Self { + data, + ..Default::default() + } + } + + /// Return a handle to the index counter for this stream + pub fn index(&self) -> BatchIndex { + self.index.clone() + } +} + +impl Stream for TestStream { + type Item = Result<RecordBatch>; + + fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> { + let next_batch = self.index.value(); + + Poll::Ready(if next_batch < self.data.len() { + let next_batch = self.index.value(); + self.index.incr(); + Some(Ok(self.data[next_batch].clone())) + } else { + None + }) + } + + fn size_hint(&self) -> (usize, Option<usize>) { + (self.data.len(), Some(self.data.len())) + } +} + +impl RecordBatchStream for TestStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + self.data[0].schema() + } +} + +/// A Mock ExecutionPlan that can be used for writing tests of other +/// ExecutionPlans +#[derive(Debug)] +pub struct MockExec { + /// the results to send back + data: Vec<Result<RecordBatch>>, + schema: SchemaRef, + /// if true (the default), sends data using a separate task to ensure the + /// batches are not available without this stream yielding first + use_task: bool, + /// if true, report unknown statistics instead of deriving them from + /// `data` (which propagates any planted errors at planning time) + unknown_statistics: bool, + cache: Arc<PlanProperties>, +} + +impl MockExec { + /// Create a new `MockExec` with a single partition that returns + /// the specified `Results`s. + /// + /// By default, the batches are not produced immediately (the + /// caller has to actually yield and another task must run) to + /// ensure any poll loops are correct. This behavior can be + /// changed with `with_use_task` + pub fn new(data: Vec<Result<RecordBatch>>, schema: SchemaRef) -> Self { + let cache = Self::compute_properties(Arc::clone(&schema)); + Self { + data, + schema, + use_task: true, + unknown_statistics: false, + cache: Arc::new(cache), + } + } + + /// If `use_task` is true (the default) then the batches are sent + /// back using a separate task to ensure the underlying stream is + /// not immediately ready + pub fn with_use_task(mut self, use_task: bool) -> Self { + self.use_task = use_task; + self + } + + /// Report unknown statistics rather than computing them from `data`. + /// + /// By default statistics are derived from `data`, which propagates any + /// planted errors when statistics are requested during planning (for + /// example when a parent node computes its properties). Use this when a + /// planted error should only surface at execution time. + pub fn with_unknown_statistics(mut self) -> Self { + self.unknown_statistics = true; + self + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for MockExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "MockExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for MockExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + /// Returns a stream which yields data + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + assert_eq!(partition, 0); + + // Result doesn't implement clone, so do it ourself + let data: Vec<_> = self + .data + .iter() + .map(|r| match r { + Ok(batch) => Ok(batch.clone()), + Err(e) => Err(clone_error(e)), + }) + .collect(); + + if self.use_task { + let mut builder = RecordBatchReceiverStream::builder(self.schema(), 2); + // send data in order but in a separate task (to ensure + // the batches are not available without the stream + // yielding). + let tx = builder.tx(); + builder.spawn(async move { + for batch in data { + println!("Sending batch via delayed stream"); + if let Err(e) = tx.send(batch).await { + println!("ERROR batch via delayed stream: {e}"); + } + } + + Ok(()) + }); + // returned stream simply reads off the rx stream + Ok(builder.build()) + } else { + // make an input that will error + let stream = futures::stream::iter(data); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } + } + + // Errors if one of the batches is an error, unless + // `with_unknown_statistics` was used + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if self.unknown_statistics || args.partition().is_some() { + return Ok(Arc::new(Statistics::new_unknown(&self.schema))); + } + let data: Result<Vec<_>> = self + .data + .iter() + .map(|r| match r { + Ok(batch) => Ok(batch.clone()), + Err(e) => Err(clone_error(e)), + }) + .collect(); + + let data = data?; + + Ok(Arc::new(common::compute_record_batch_statistics( + &[data], + &self.schema, + None, + ))) + } +} + +fn clone_error(e: &DataFusionError) -> DataFusionError { + use DataFusionError::*; + match e { + Execution(msg) => Execution(msg.to_string()), + _ => unimplemented!(), + } +} + +/// A Mock ExecutionPlan that does not start producing input until a +/// barrier is called +#[derive(Debug)] +pub struct BarrierExec { + /// partitions to send back + data: Vec<Vec<RecordBatch>>, + schema: SchemaRef, + + /// all streams wait on this barrier to produce + start_data_barrier: Option<Arc<Barrier>>, + + /// the stream wait for this to return Poll::Ready(None) + finish_barrier: Option<Arc<(Barrier, AtomicUsize)>>, + + cache: Arc<PlanProperties>, + + log: bool, +} + +impl BarrierExec { + /// Create a new exec with some number of partitions. + pub fn new(data: Vec<Vec<RecordBatch>>, schema: SchemaRef) -> Self { + // wait for all streams and the input + let barrier = Some(Arc::new(Barrier::new(data.len() + 1))); + let cache = Self::compute_properties(Arc::clone(&schema), &data); + Self { + data, + schema, + start_data_barrier: barrier, + cache: Arc::new(cache), + finish_barrier: None, + log: true, + } + } + + pub fn with_log(mut self, log: bool) -> Self { + self.log = log; + self + } + + pub fn without_start_barrier(mut self) -> Self { + self.start_data_barrier = None; + self + } + + pub fn with_finish_barrier(mut self) -> Self { + let barrier = Arc::new(( + // wait for all streams and the input + Barrier::new(self.data.len() + 1), + AtomicUsize::new(0), + )); + + self.finish_barrier = Some(barrier); + self + } + + /// wait until all the input streams and this function is ready + pub async fn wait(&self) { + let barrier = &self + .start_data_barrier + .as_ref() + .expect("Must only be called when having a start barrier"); + if self.log { + println!("BarrierExec::wait waiting on barrier"); + } + barrier.wait().await; + if self.log { + println!("BarrierExec::wait done waiting"); + } + } + + pub async fn wait_finish(&self) { + let (barrier, _) = &self + .finish_barrier + .as_deref() + .expect("Must only be called when having a finish barrier"); + + if self.log { + println!("BarrierExec::wait_finish waiting on barrier"); + } + barrier.wait().await; + if self.log { + println!("BarrierExec::wait_finish done waiting"); + } + } + + /// Return true if the finish barrier has been reached in all partitions + pub fn is_finish_barrier_reached(&self) -> bool { + let (_, reached_finish) = self + .finish_barrier + .as_deref() + .expect("Must only be called when having finish barrier"); + + reached_finish.load(Ordering::Relaxed) == self.data.len() + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + data: &[Vec<RecordBatch>], + ) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(data.len()), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for BarrierExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "BarrierExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for BarrierExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + /// Returns a stream which yields data + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + assert!(partition < self.data.len()); + + let mut builder = RecordBatchReceiverStream::builder(self.schema(), 2); + + // task simply sends data in order after barrier is reached + let data = self.data[partition].clone(); + let start_barrier = self.start_data_barrier.as_ref().map(Arc::clone); + let finish_barrier = self.finish_barrier.as_ref().map(Arc::clone); + let log = self.log; + let tx = builder.tx(); + builder.spawn(async move { + if let Some(barrier) = start_barrier { + if log { + println!("Partition {partition} waiting on barrier"); + } + barrier.wait().await; + } + for batch in data { + if log { + println!("Partition {partition} sending batch"); + } + if let Err(e) = tx.send(Ok(batch)).await { + println!("ERROR batch via barrier stream stream: {e}"); + } + } + if let Some((barrier, reached_finish)) = finish_barrier.as_deref() { + if log { + println!("Partition {partition} waiting on finish barrier"); + } + reached_finish.fetch_add(1, Ordering::Relaxed); + barrier.wait().await; + } + + Ok(()) + }); + + // returned stream simply reads off the rx stream + Ok(builder.build()) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if args.partition().is_some() { + return Ok(Arc::new(Statistics::new_unknown(&self.schema))); + } + Ok(Arc::new(common::compute_record_batch_statistics( + &self.data, + &self.schema, + None, + ))) + } +} + +/// A mock execution plan that errors on a call to execute +#[derive(Debug)] +pub struct ErrorExec { + cache: Arc<PlanProperties>, +} + +impl Default for ErrorExec { + fn default() -> Self { + Self::new() + } +} + +impl ErrorExec { + pub fn new() -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + "dummy", + DataType::Int64, + true, + )])); + let cache = Self::compute_properties(schema); + Self { + cache: Arc::new(cache), + } + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for ErrorExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ErrorExec") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for ErrorExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + unimplemented!() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + /// Returns a stream which yields data + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + internal_err!("ErrorExec, unsurprisingly, errored in partition {partition}") + } +} + +/// A mock execution plan that simply returns the provided statistics +#[derive(Debug, Clone)] +pub struct StatisticsExec { + stats: Statistics, + schema: Arc<Schema>, + cache: Arc<PlanProperties>, +} +impl StatisticsExec { + pub fn new(stats: Statistics, schema: Schema) -> Self { + assert_eq!( + stats.column_statistics.len(), + schema.fields().len(), + "if defined, the column statistics vector length should be the number of fields" + ); + let cache = Self::compute_properties(Arc::new(schema.clone())); + Self { + stats, + schema: Arc::new(schema), + cache: Arc::new(cache), + } + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(2), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for StatisticsExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "StatisticsExec: col_count={}, row_count={:?}", + self.schema.fields().len(), + self.stats.num_rows, + ) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for StatisticsExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + unimplemented!("This plan only serves for testing statistics") + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::new(if args.partition().is_some() { + Statistics::new_unknown(&self.schema) + } else { + self.stats.clone() + })) + } +} + +/// Execution plan that emits streams that block forever. +/// +/// This is useful to test shutdown / cancellation behavior of certain execution plans. +#[derive(Debug)] +pub struct BlockingExec { + /// Schema that is mocked by this plan. + schema: SchemaRef, + + /// Ref-counting helper to check if the plan and the produced stream are still in memory. + refs: Arc<()>, + cache: Arc<PlanProperties>, +} + +impl BlockingExec { + /// Create new [`BlockingExec`] with a give schema and number of partitions. + pub fn new(schema: SchemaRef, n_partitions: usize) -> Self { + let cache = Self::compute_properties(Arc::clone(&schema), n_partitions); + Self { + schema, + refs: Default::default(), + cache: Arc::new(cache), + } + } + + /// Weak pointer that can be used for ref-counting this execution plan and its streams. + /// + /// Use [`Weak::strong_count`] to determine if the plan itself and its streams are dropped (should be 0 in that + /// case). Note that tokio might take some time to cancel spawned tasks, so you need to wrap this check into a retry + /// loop. Use [`assert_strong_count_converges_to_zero`] to archive this. + pub fn refs(&self) -> Weak<()> { + Arc::downgrade(&self.refs) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef, n_partitions: usize) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(n_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for BlockingExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "BlockingExec",) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for BlockingExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + // this is a leaf node and has no children + vec![] + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + internal_err!("Children cannot be replaced in {self:?}") + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(BlockingStream { + schema: Arc::clone(&self.schema), + _refs: Arc::clone(&self.refs), + })) + } +} + +/// A [`RecordBatchStream`] that is pending forever. +#[derive(Debug)] +pub struct BlockingStream { + /// Schema mocked by this stream. + schema: SchemaRef, + + /// Ref-counting helper to check if the stream are still in memory. + _refs: Arc<()>, +} + +impl Stream for BlockingStream { + type Item = Result<RecordBatch>; + + fn poll_next( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + Poll::Pending + } +} + +impl RecordBatchStream for BlockingStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Asserts that the strong count of the given [`Weak`] pointer converges to zero. +/// +/// This might take a while but has a timeout. +pub async fn assert_strong_count_converges_to_zero<T>(refs: Weak<T>) { + tokio::time::timeout(std::time::Duration::from_secs(10), async { + loop { + if dbg!(Weak::strong_count(&refs)) == 0 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); +} + +/// Execution plan that emits streams that panics. +/// +/// This is useful to test panic handling of certain execution plans. +#[derive(Debug)] +pub struct PanicExec { + /// Schema that is mocked by this plan. + schema: SchemaRef, + + /// Number of output partitions. Each partition will produce this + /// many empty output record batches prior to panicking + batches_until_panics: Vec<usize>, + cache: Arc<PlanProperties>, +} + +impl PanicExec { + /// Create new [`PanicExec`] with a give schema and number of + /// partitions, which will each panic immediately. + pub fn new(schema: SchemaRef, n_partitions: usize) -> Self { + let batches_until_panics = vec![0; n_partitions]; + let cache = Self::compute_properties(Arc::clone(&schema), &batches_until_panics); + Self { + schema, + batches_until_panics, + cache: Arc::new(cache), + } + } + + /// Set the number of batches prior to panic for a partition + pub fn with_partition_panic(mut self, partition: usize, count: usize) -> Self { + self.batches_until_panics[partition] = count; + self + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + batches_until_panics: &[usize], + ) -> PlanProperties { + let num_partitions = batches_until_panics.len(); + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +impl DisplayAs for PanicExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "PanicExec",) + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ExecutionPlan for PanicExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + // this is a leaf node and has no children + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + internal_err!("Children cannot be replaced in {:?}", self) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + Ok(Box::pin(PanicStream { + partition, + batches_until_panic: self.batches_until_panics[partition], + schema: Arc::clone(&self.schema), + ready: false, + })) + } +} + +/// A [`RecordBatchStream`] that yields every other batch and panics +/// after `batches_until_panic` batches have been produced. +/// +/// Useful for testing the behavior of streams on panic +#[derive(Debug)] +struct PanicStream { + /// Which partition was this + partition: usize, + /// How may batches will be produced until panic + batches_until_panic: usize, + /// Schema mocked by this stream. + schema: SchemaRef, + /// Should we return ready ? + ready: bool, +} + +impl Stream for PanicStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + if self.batches_until_panic > 0 { + if self.ready { + self.batches_until_panic -= 1; + self.ready = false; + let batch = RecordBatch::new_empty(Arc::clone(&self.schema)); + return Poll::Ready(Some(Ok(batch))); + } else { + self.ready = true; + // get called again + cx.waker().wake_by_ref(); + return Poll::Pending; + } + } + panic!("PanickingStream did panic: {}", self.partition) + } +} + +impl RecordBatchStream for PanicStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/topk/mod.rs b/native/vendor/datafusion-physical-plan/src/topk/mod.rs new file mode 100644 index 0000000..1e3efff --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/topk/mod.rs
@@ -0,0 +1,3169 @@ +// 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. + +//! TopK: Combination of Sort / LIMIT + +use arrow::{ + array::{Array, AsArray}, + compute::{ + BatchCoalescer, FilterBuilder, interleave_record_batch, prep_null_mask_filter, + take_record_batch, + }, + row::{OwnedRow, RowConverter, Rows, SortField}, +}; +use datafusion_expr::{ColumnarValue, Operator}; +use std::mem::size_of; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; +use std::{cmp::Ordering, collections::BinaryHeap, sync::Arc}; + +use super::metrics::{ + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + RecordOutput, +}; +use crate::spill::get_record_batch_memory_size; +use crate::{SendableRecordBatchStream, stream::RecordBatchStreamAdapter}; + +use arrow::array::{ArrayRef, RecordBatch, UInt32Array}; +use arrow::datatypes::SchemaRef; +use datafusion_common::{ + HashMap, Result, ScalarValue, internal_datafusion_err, internal_err, +}; +use datafusion_execution::{ + memory_pool::{MemoryConsumer, MemoryReservation}, + runtime_env::RuntimeEnv, +}; +use datafusion_physical_expr::{ + PhysicalExpr, + expressions::{BinaryExpr, DynamicFilterPhysicalExpr, is_not_null, is_null, lit}, +}; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +use parking_lot::RwLock; + +/// TopK +/// +/// # Background +/// +/// "Top K" is a common query optimization used for queries such as +/// "find the top 3 customers by revenue". The (simplified) SQL for +/// such a query might be: +/// +/// ```sql +/// SELECT customer_id, revenue FROM 'sales.csv' ORDER BY revenue DESC limit 3; +/// ``` +/// +/// The simple plan would be: +/// +/// ```sql +/// > explain SELECT customer_id, revenue FROM sales ORDER BY revenue DESC limit 3; +/// +--------------+----------------------------------------+ +/// | plan_type | plan | +/// +--------------+----------------------------------------+ +/// | logical_plan | Limit: 3 | +/// | | Sort: revenue DESC NULLS FIRST | +/// | | Projection: customer_id, revenue | +/// | | TableScan: sales | +/// +--------------+----------------------------------------+ +/// ``` +/// +/// While this plan produces the correct answer, it will fully sorts the +/// input before discarding everything other than the top 3 elements. +/// +/// The same answer can be produced by simply keeping track of the top +/// K=3 elements, reducing the total amount of required buffer memory. +/// +/// # Partial Sort Optimization +/// +/// This implementation additionally optimizes queries where the input is already +/// partially sorted by a common prefix of the requested ordering. If subsequent +/// rows are guaranteed to be strictly greater (in sort order) than a known TopK +/// boundary on this prefix, the operator safely terminates early. +/// +/// For a local TopK, that boundary comes from the local heap once it has K rows. +/// For a partitioned `SortExec`, a shared dynamic-filter threshold can provide +/// the same prefix boundary before a lagging partition has filled its local heap. +/// +/// ## Example +/// +/// For input sorted by `(day DESC)`, but not by `timestamp`, a query such as: +/// +/// ```sql +/// SELECT day, timestamp FROM sensor ORDER BY day DESC, timestamp DESC LIMIT 10; +/// ``` +/// +/// can terminate scanning early once sufficient rows from the latest days have been +/// collected, skipping older data. +/// +/// # Structure +/// +/// This operator tracks the top K items using a `TopKHeap`. +pub struct TopK { + /// schema of the output (and the input) + schema: SchemaRef, + /// Runtime metrics + metrics: TopKMetrics, + /// Reservation + reservation: MemoryReservation, + /// The target number of rows for output batches + batch_size: usize, + /// sort expressions + expr: LexOrdering, + /// row converter, for sort keys + row_converter: RowConverter, + /// scratch space for converting rows + scratch_rows: Rows, + /// stores the top k values and their sort key values, in order + heap: TopKHeap, + /// row converter, for common keys between the sort keys and the input ordering + common_sort_prefix_converter: Option<RowConverter>, + /// Common sort prefix between the input and the sort expressions to allow early exit optimization + common_sort_prefix: Arc<[PhysicalSortExpr]>, + /// Filter matching the state of the `TopK` heap used for dynamic filter pushdown + filter: Arc<RwLock<TopKDynamicFilters>>, + /// If true, indicates that all rows of subsequent batches are guaranteed + /// to be greater (by byte order, after row conversion) than the top K, + /// which means the top K won't change and the computation can be finished early. + pub(crate) finished: bool, +} + +/// For more background, please also see the [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog] +/// +/// [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog]: https://datafusion.apache.org/blog/2025/09/10/dynamic-filters +#[derive(Debug)] +pub struct TopKDynamicFilters { + /// The current threshold shared by all TopK emitters that use this dynamic + /// filter. Any emitter may tighten it. + /// + /// The full sort-key row and common-prefix row are stored together so they + /// always describe the same heap row. + shared_threshold: Option<TopKThreshold>, + /// The expression used to evaluate the dynamic filter + /// Only updated when lock held for the duration of the update + expr: Arc<DynamicFilterPhysicalExpr>, + /// Number of local TopK emitters that have not called `emit` yet. + /// + /// A partition-preserving `SortExec` creates one local TopK per output + /// partition. The shared dynamic filter is complete only after every local + /// TopK has emitted. + /// + /// `emit` only needs a read guard on the shared filter wrapper, so + /// concurrent emitters use this atomic counter instead of taking an + /// exclusive lock just to mark their partition done. + remaining_topk_emitters: AtomicUsize, +} + +#[derive(Debug, Clone)] +struct TopKThreshold { + /// The full sort-key row bytes for efficient comparison. + full_sort_key_row: Vec<u8>, + /// The same heap row encoded with the common-prefix converter, when the + /// input ordering shares a prefix with the TopK ordering. + /// + /// This lets each partition stop from a shared TopK threshold even if its + /// local heap has not filled yet. + common_prefix_row: Option<Vec<u8>>, +} + +impl TopKThreshold { + fn new(full_sort_key_row: Vec<u8>, common_prefix_row: Option<Vec<u8>>) -> Self { + Self { + full_sort_key_row, + common_prefix_row, + } + } + + fn full_sort_key_row(&self) -> &[u8] { + self.full_sort_key_row.as_slice() + } + + fn common_prefix_row(&self) -> Option<&[u8]> { + self.common_prefix_row.as_deref() + } + + fn is_more_selective_than(&self, current: &Self) -> bool { + self.full_sort_key_row() < current.full_sort_key_row() + } +} + +#[derive(Clone, Copy)] +struct TopKHeapBoundaryRow<'a> { + row: &'a TopKRow, +} + +impl<'a> TopKHeapBoundaryRow<'a> { + fn new(row: &'a TopKRow) -> Self { + Self { row } + } + + fn full_sort_key_row(&self) -> &[u8] { + self.row.row() + } + + fn is_more_selective_than(&self, current: Option<&TopKThreshold>) -> bool { + current + .map(|current| self.full_sort_key_row() < current.full_sort_key_row()) + .unwrap_or(true) + } +} + +#[derive(Clone, Copy)] +struct TopKHeapBoundary<'a> { + row: &'a TopKRow, + batch: &'a RecordBatch, +} + +impl<'a> TopKHeapBoundary<'a> { + fn new(row: &'a TopKRow, batch: &'a RecordBatch) -> Self { + Self { row, batch } + } + + fn threshold_values( + &self, + sort_exprs: &[PhysicalSortExpr], + ) -> Result<Vec<ScalarValue>> { + let mut scalar_values = Vec::with_capacity(sort_exprs.len()); + for sort_expr in sort_exprs { + let value = sort_expr + .expr + .evaluate(&self.batch.slice(self.row.index, 1))?; + + let scalar = match value { + ColumnarValue::Scalar(scalar) => scalar, + ColumnarValue::Array(array) if array.len() == 1 => { + ScalarValue::try_from_array(&array, 0)? + } + array => { + return internal_err!("Expected a scalar value, got {:?}", array); + } + }; + scalar_values.push(scalar); + } + + Ok(scalar_values) + } + + fn threshold(&self, common_prefix_row: Option<Vec<u8>>) -> TopKThreshold { + TopKThreshold::new(self.row.row().to_vec(), common_prefix_row) + } +} + +impl TopKDynamicFilters { + /// Create a new `TopKDynamicFilters` with the given expression + pub fn new(expr: Arc<DynamicFilterPhysicalExpr>) -> Self { + Self::new_with_topk_emitter_count(expr, 1) + } + + /// Create a new `TopKDynamicFilters` with the expected number of local + /// TopK emitters that share it. + pub fn new_with_topk_emitter_count( + expr: Arc<DynamicFilterPhysicalExpr>, + topk_emitter_count: usize, + ) -> Self { + debug_assert!(topk_emitter_count > 0); + Self { + shared_threshold: None, + expr, + remaining_topk_emitters: AtomicUsize::new(topk_emitter_count), + } + } + + pub fn expr(&self) -> Arc<DynamicFilterPhysicalExpr> { + Arc::clone(&self.expr) + } + + fn mark_topk_emitted(&self) { + let previous = self + .remaining_topk_emitters + .fetch_update( + AtomicOrdering::AcqRel, + AtomicOrdering::Acquire, + |remaining| remaining.checked_sub(1), + ) + .unwrap_or(0); + debug_assert!( + previous > 0, + "TopK dynamic filter emitter completed more times than expected" + ); + + if previous == 1 { + self.expr.mark_complete(); + } + } +} + +// Guesstimate for memory allocation: estimated number of bytes used per row in the RowConverter +const ESTIMATED_BYTES_PER_ROW: usize = 20; + +/// Owned data of a row that was just evicted from a [`TopKHeap`]. +/// +/// Returned by [`TopKHeap::add`] so that callers (e.g. rank-aware +/// wrappers that retain boundary ties) can decide whether to retain +/// the evicted row externally. The underlying batch is captured +/// before the heap's internal `RecordBatchStore` decrements the +/// batch's use count, so the data remains accessible even if the +/// heap drops its internal reference to the batch. +#[derive(Debug, Clone)] +pub(crate) struct EvictedRow { + /// The record batch the evicted row came from. + pub batch: RecordBatch, + /// Row index within `batch`. + pub index: usize, + /// Encoded ORDER BY tuple for the evicted row, in [`arrow::row`] format. + pub row_bytes: Vec<u8>, +} + +pub(crate) fn build_sort_fields( + ordering: &[PhysicalSortExpr], + schema: &SchemaRef, +) -> Result<Vec<SortField>> { + ordering + .iter() + .map(|e| { + Ok(SortField::new_with_options( + e.expr.data_type(schema)?, + e.options, + )) + }) + .collect::<Result<_>>() +} + +impl TopK { + /// Create a new [`TopK`] that stores the top `k` values, as + /// defined by the sort expressions in `expr`. + // TODO: make a builder or some other nicer API + #[expect(clippy::too_many_arguments)] + #[expect(clippy::needless_pass_by_value)] + pub fn try_new( + partition_id: usize, + schema: SchemaRef, + common_sort_prefix: Vec<PhysicalSortExpr>, + expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: Arc<RuntimeEnv>, + metrics: &ExecutionPlanMetricsSet, + filter: Arc<RwLock<TopKDynamicFilters>>, + ) -> Result<Self> { + let reservation = MemoryConsumer::new(format!("TopK[{partition_id}]")) + .register(&runtime.memory_pool); + + let sort_fields = build_sort_fields(&expr, &schema)?; + + // TODO there is potential to add special cases for single column sort fields + // to improve performance + let row_converter = RowConverter::new(sort_fields)?; + let scratch_rows = + row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + let common_prefix_row_converter = if common_sort_prefix.is_empty() { + None + } else { + let input_sort_fields = build_sort_fields(&common_sort_prefix, &schema)?; + Some(RowConverter::new(input_sort_fields)?) + }; + + Ok(Self { + schema: Arc::clone(&schema), + metrics: TopKMetrics::new(metrics, partition_id), + reservation, + batch_size, + expr, + row_converter, + scratch_rows, + heap: TopKHeap::new(k), + common_sort_prefix_converter: common_prefix_row_converter, + common_sort_prefix: Arc::from(common_sort_prefix), + finished: false, + filter, + }) + } + + /// Insert `batch`, remembering if any of its values are among + /// the top k seen so far. + #[expect(clippy::needless_pass_by_value)] + pub fn insert_batch(&mut self, batch: RecordBatch) -> Result<()> { + // Updates on drop + let baseline = self.metrics.baseline.clone(); + let _timer = baseline.elapsed_compute().timer(); + + let mut sort_keys: Vec<ArrayRef> = self + .expr + .iter() + .map(|expr| { + let value = expr.expr.evaluate(&batch)?; + value.into_array(batch.num_rows()) + }) + .collect::<Result<Vec<_>>>()?; + + let mut selected_rows = None; + + // If a filter is provided, update it with the new rows + let filter = self.filter.read().expr.current()?; + let filtered = filter.evaluate(&batch)?; + let num_rows = batch.num_rows(); + let array = filtered.into_array(num_rows)?; + let mut filter = array.as_boolean().clone(); + if !filter.has_true() { + // The heap is unchanged, but a fully rejected batch can still prove + // that the shared sort prefix has passed the heap boundary. + self.attempt_early_completion(&batch)?; + return Ok(()); + } + // only update the keys / rows if the filter does not match all rows + if filter.null_count() > 0 || filter.has_false() { + // Indices in `set_indices` should be correct if filter contains nulls + // So we prepare the filter here. Note this is also done in the `FilterBuilder` + // so there is no overhead to do this here. + if filter.nulls().is_some() { + filter = prep_null_mask_filter(&filter); + } + + let filter_predicate = FilterBuilder::new(&filter); + let filter_predicate = if sort_keys.len() > 1 { + // Optimize filter when it has multiple sort keys + filter_predicate.optimize().build() + } else { + filter_predicate.build() + }; + selected_rows = Some(filter); + sort_keys = sort_keys + .iter() + .map(|key| filter_predicate.filter(key).map_err(|x| x.into())) + .collect::<Result<Vec<_>>>()?; + } + // reuse existing `Rows` to avoid reallocations + let rows = &mut self.scratch_rows; + rows.clear(); + self.row_converter.append(rows, &sort_keys)?; + + let mut batch_entry = self.heap.register_batch(batch.clone()); + + let replacements = match selected_rows { + Some(filter) => { + self.find_new_topk_items(filter.values().set_indices(), &mut batch_entry) + } + None => self.find_new_topk_items(0..sort_keys[0].len(), &mut batch_entry), + }; + + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + + self.heap.insert_batch_entry(batch_entry); + + // conserve memory + self.heap.maybe_compact()?; + + // update memory reservation + self.reservation.try_resize(self.size())?; + + // flag the topK as finished if we know that all + // subsequent batches are guaranteed to be greater (by byte order, after row conversion) than the top K, + // which means the top K won't change and the computation can be finished early. + self.attempt_early_completion(&batch)?; + + // update the filter representation of our TopK heap + self.update_filter()?; + } else { + // The heap did not change, but this batch's prefix may still prove + // that no later rows can enter the TopK. + self.attempt_early_completion(&batch)?; + } + + Ok(()) + } + + fn find_new_topk_items( + &mut self, + items: impl Iterator<Item = usize>, + batch_entry: &mut RecordBatchEntry, + ) -> usize { + let mut replacements = 0; + let rows = &mut self.scratch_rows; + for (index, row) in items.zip(rows.iter()) { + match self.heap.max() { + // heap has k items, and the new row is greater than the + // current max in the heap ==> it is not a new topk + Some(max_row) if row.as_ref() >= max_row.row() => {} + // don't yet have k items or new item is lower than the currently k low values + None | Some(_) => { + self.heap.add(batch_entry, row, index); + replacements += 1; + } + } + } + replacements + } + + fn current_heap_boundary_row(&self) -> Option<TopKHeapBoundaryRow<'_>> { + self.heap.max().map(TopKHeapBoundaryRow::new) + } + + fn current_heap_boundary(&self) -> Result<Option<TopKHeapBoundary<'_>>> { + let Some(row) = self.heap.max() else { + return Ok(None); + }; + + self.heap_boundary(row).map(Some) + } + + fn heap_boundary<'a>(&'a self, row: &'a TopKRow) -> Result<TopKHeapBoundary<'a>> { + let batch_entry = self + .heap + .store + .get(row.batch_id) + .ok_or_else(|| internal_datafusion_err!("Invalid batch ID in TopKRow"))?; + + Ok(TopKHeapBoundary::new(row, &batch_entry.batch)) + } + + /// Update the filter representation of our TopK heap. + /// For example, given the sort expression `ORDER BY a DESC, b ASC LIMIT 3`, + /// and the current heap values `[(1, 5), (1, 4), (2, 3)]`, + /// the filter will be updated to: + /// + /// ```sql + /// (a > 1 OR (a = 1 AND b < 5)) AND + /// (a > 1 OR (a = 1 AND b < 4)) AND + /// (a > 2 OR (a = 2 AND b < 3)) + /// ``` + fn update_filter(&mut self) -> Result<()> { + // If the heap doesn't have k elements yet, we can't create thresholds + let Some(boundary_row) = self.current_heap_boundary_row() else { + return Ok(()); + }; + + // Fast path: check if the current value in topk is better than what is + // currently set in the filter with a read only lock + let needs_update = { + let filter = self.filter.read(); + boundary_row.is_more_selective_than(filter.shared_threshold.as_ref()) + }; + + // exit early if the current values are better + if !needs_update { + return Ok(()); + } + + let boundary = self.heap_boundary(boundary_row.row)?; + + // Extract scalar values BEFORE acquiring lock to reduce critical section + let thresholds = boundary.threshold_values(&self.expr)?; + + // Build the filter expression OUTSIDE any synchronization + let predicate = Self::build_filter_expression(&self.expr, &thresholds)?; + let new_threshold = + boundary.threshold(self.encode_topk_common_prefix_row(boundary)?); + + // update the threshold. Since there was a lock gap, we must check if it is still the best + // may have changed while we were building the expression without the lock + let mut filter = self.filter.write(); + let still_needs_update = filter + .shared_threshold + .as_ref() + .map(|current| new_threshold.is_more_selective_than(current)) + .unwrap_or(true); + if !still_needs_update { + // some other thread updated the threshold to a better one while we + // were building so there is no need to update the filter + return Ok(()); + } + filter.shared_threshold = Some(new_threshold); + + // Update the filter expression + if let Some(pred) = predicate + && !pred.eq(&lit(true)) + { + filter.expr.update(pred)?; + } + + Ok(()) + } + + /// Build the filter expression with the given thresholds. + /// This is now called outside of any locks to reduce critical section time. + fn build_filter_expression( + sort_exprs: &[PhysicalSortExpr], + thresholds: &[ScalarValue], + ) -> Result<Option<Arc<dyn PhysicalExpr>>> { + // Create filter expressions for each threshold + let mut filters: Vec<Arc<dyn PhysicalExpr>> = + Vec::with_capacity(thresholds.len()); + + let mut prev_sort_expr: Option<Arc<dyn PhysicalExpr>> = None; + for (sort_expr, value) in sort_exprs.iter().zip(thresholds.iter()) { + // Create the appropriate operator based on sort order + let op = if sort_expr.options.descending { + // For descending sort, we want col > threshold (exclude smaller values) + Operator::Gt + } else { + // For ascending sort, we want col < threshold (exclude larger values) + Operator::Lt + }; + + let value_null = value.is_null(); + + let comparison = Arc::new(BinaryExpr::new( + Arc::clone(&sort_expr.expr), + op, + lit(value.clone()), + )); + + let comparison_with_null = match (sort_expr.options.nulls_first, value_null) { + // For nulls first, transform to (threshold.value is not null) and (threshold.expr is null or comparison) + (true, true) => lit(false), + (true, false) => Arc::new(BinaryExpr::new( + is_null(Arc::clone(&sort_expr.expr))?, + Operator::Or, + comparison, + )), + // For nulls last, transform to (threshold.value is null and threshold.expr is not null) + // or (threshold.value is not null and comparison) + (false, true) => is_not_null(Arc::clone(&sort_expr.expr))?, + (false, false) => comparison, + }; + + let mut eq_expr = Arc::new(BinaryExpr::new( + Arc::clone(&sort_expr.expr), + Operator::Eq, + lit(value.clone()), + )); + + if value_null { + eq_expr = Arc::new(BinaryExpr::new( + is_null(Arc::clone(&sort_expr.expr))?, + Operator::Or, + eq_expr, + )); + } + + // For a query like order by a, b, the filter for column `b` is only applied if + // the condition a = threshold.value (considering null equality) is met. + // Therefore, we add equality predicates for all preceding fields to the filter logic of the current field, + // and include the current field's equality predicate in `prev_sort_expr` for use with subsequent fields. + match prev_sort_expr.take() { + None => { + prev_sort_expr = Some(eq_expr); + filters.push(comparison_with_null); + } + Some(p) => { + filters.push(Arc::new(BinaryExpr::new( + Arc::clone(&p), + Operator::And, + comparison_with_null, + ))); + + prev_sort_expr = + Some(Arc::new(BinaryExpr::new(p, Operator::And, eq_expr))); + } + } + } + + let dynamic_predicate = filters + .into_iter() + .reduce(|a, b| Arc::new(BinaryExpr::new(a, Operator::Or, b))); + + Ok(dynamic_predicate) + } + + /// If input ordering shares a common sort prefix with the TopK, + /// check if the computation can be finished early. + /// + /// This is the case if the last row of the current batch is strictly + /// greater than either the shared dynamic-filter threshold prefix or the max + /// row in the local heap, comparing only on the shared prefix columns. + fn attempt_early_completion(&mut self, batch: &RecordBatch) -> Result<()> { + // Early exit if the batch is empty as there is no last row to extract from it. + if batch.num_rows() == 0 { + return Ok(()); + } + + // common_prefix_row_converter is only `Some` if the input ordering has a common prefix with the TopK, + // so early exit if it is `None`. + let Some(prefix_converter) = &self.common_sort_prefix_converter else { + return Ok(()); + }; + + // Evaluate the prefix for the last row of the current batch. + let last_row_idx = batch.num_rows() - 1; + let mut batch_prefix_scratch = + prefix_converter.empty_rows(1, ESTIMATED_BYTES_PER_ROW); // 1 row with capacity ESTIMATED_BYTES_PER_ROW + + self.append_common_prefix_row( + prefix_converter, + batch, + last_row_idx, + &mut batch_prefix_scratch, + )?; + let batch_common_prefix_row = batch_prefix_scratch.row(0); + let batch_common_prefix = batch_common_prefix_row.as_ref(); + + let finished_by_shared_threshold = self + .filter + .read() + .shared_threshold + .as_ref() + .and_then(TopKThreshold::common_prefix_row) + .map(|common_prefix_row| batch_common_prefix > common_prefix_row) + .unwrap_or(false); + if finished_by_shared_threshold { + self.finished = true; + return Ok(()); + } + + // Early exit only from the local heap once it has a full boundary row. + let Some(boundary) = self.current_heap_boundary()? else { + return Ok(()); + }; + + if self.batch_prefix_exceeds_heap_boundary(batch_common_prefix, boundary)? { + self.finished = true; + } + + Ok(()) + } + + fn batch_prefix_exceeds_heap_boundary( + &self, + batch_common_prefix: &[u8], + boundary: TopKHeapBoundary<'_>, + ) -> Result<bool> { + let Some(heap_common_prefix_row) = + self.encode_topk_common_prefix_row(boundary)? + else { + return Ok(false); + }; + + Ok(batch_common_prefix > heap_common_prefix_row.as_slice()) + } + + fn encode_topk_common_prefix_row( + &self, + boundary: TopKHeapBoundary<'_>, + ) -> Result<Option<Vec<u8>>> { + let Some(prefix_converter) = &self.common_sort_prefix_converter else { + return Ok(None); + }; + + let mut scratch = prefix_converter.empty_rows(1, ESTIMATED_BYTES_PER_ROW); + self.append_common_prefix_row( + prefix_converter, + boundary.batch, + boundary.row.index, + &mut scratch, + )?; + Ok(Some(scratch.row(0).as_ref().to_vec())) + } + + fn append_common_prefix_row( + &self, + prefix_converter: &RowConverter, + batch: &RecordBatch, + row_idx: usize, + scratch: &mut Rows, + ) -> Result<()> { + let row = batch.slice(row_idx, 1); + let prefix_columns: Vec<ArrayRef> = self + .common_sort_prefix + .iter() + .map(|expr| expr.expr.evaluate(&row)?.into_array(1)) + .collect::<Result<_>>()?; + + prefix_converter.append(scratch, &prefix_columns)?; + Ok(()) + } + + /// Returns the top k results broken into `batch_size` [`RecordBatch`]es, consuming the heap + pub fn emit(self) -> Result<SendableRecordBatchStream> { + let Self { + schema, + metrics, + reservation: _, + batch_size, + expr: _, + row_converter: _, + scratch_rows: _, + mut heap, + common_sort_prefix_converter: _, + common_sort_prefix: _, + finished: _, + filter, + } = self; + let _timer = metrics.baseline.elapsed_compute().timer(); // time updated on drop + + // Mark this local TopK as emitted. For shared filters, the final + // local emitter marks the dynamic filter complete. + filter.read().mark_topk_emitted(); + + // break into record batches as needed + let mut batches = vec![]; + if let Some(mut batch) = heap.emit()? { + (&batch).record_output(&metrics.baseline); + + loop { + if batch.num_rows() <= batch_size { + batches.push(Ok(batch)); + break; + } else { + batches.push(Ok(batch.slice(0, batch_size))); + let remaining_length = batch.num_rows() - batch_size; + batch = batch.slice(batch_size, remaining_length); + } + } + }; + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches), + ))) + } + + /// return the size of memory used by this operator, in bytes + fn size(&self) -> usize { + size_of::<Self>() + + self.row_converter.size() + + self.scratch_rows.size() + + self.heap.size() + } +} + +struct TopKMetrics { + /// metrics + pub baseline: BaselineMetrics, + + /// count of how many rows were replaced in the heap + pub row_replacements: Count, +} + +impl TopKMetrics { + fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + baseline: BaselineMetrics::new(metrics, partition), + row_replacements: MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("row_replacements", partition), + } + } +} + +/// This structure keeps at most the *smallest* k items, using the +/// [arrow::row] format for sort keys. While it is called "topK" for +/// values like `1, 2, 3, 4, 5` the "top 3" really means the +/// *smallest* 3 , `1, 2, 3`, not the *largest* 3 `3, 4, 5`. +/// +/// Using the `Row` format handles things such as ascending vs +/// descending and nulls first vs nulls last. +struct TopKHeap { + /// The maximum number of elements to store in this heap. + k: usize, + /// Storage for up at most `k` items using a BinaryHeap. Reversed + /// so that the smallest k so far is on the top + inner: BinaryHeap<TopKRow>, + /// Storage the original row values (TopKRow only has the sort key) + store: RecordBatchStore, + /// The size of all owned data held by this heap + owned_bytes: usize, +} + +impl TopKHeap { + fn new(k: usize) -> Self { + assert!(k > 0); + Self { + k, + inner: BinaryHeap::new(), + store: RecordBatchStore::new(), + owned_bytes: 0, + } + } + + /// Register a [`RecordBatch`] with the heap, returning the + /// appropriate entry + pub fn register_batch(&mut self, batch: RecordBatch) -> RecordBatchEntry { + self.store.register(batch) + } + + /// Insert a [`RecordBatchEntry`] created by a previous call to + /// [`Self::register_batch`] into storage. + pub fn insert_batch_entry(&mut self, entry: RecordBatchEntry) { + self.store.insert(entry) + } + + /// Returns the largest value stored by the heap if there are k + /// items, otherwise returns None. Remember this structure is + /// keeping the "smallest" k values + fn max(&self) -> Option<&TopKRow> { + if self.inner.len() < self.k { + None + } else { + self.inner.peek() + } + } + + /// Adds `row` to this heap. If inserting this new item would + /// increase the size past `k`, removes the previously smallest + /// item. + /// + /// Returns `Some(EvictedRow)` if an existing row was evicted to + /// make room for `row`, or `None` if the row was inserted into a + /// non-full heap. + fn add( + &mut self, + batch_entry: &mut RecordBatchEntry, + row: impl AsRef<[u8]>, + index: usize, + ) -> Option<EvictedRow> { + let batch_id = batch_entry.id; + batch_entry.uses += 1; + + assert!(self.inner.len() <= self.k); + let row = row.as_ref(); + + // Reuse storage for evicted item if possible + if self.inner.len() == self.k { + let mut prev_min = self.inner.peek_mut().unwrap(); + + // Capture evicted row data before `unuse` (which may GC the + // batch from the store) and `replace_with` (which overwrites + // `prev_min` in place). The batch comes from `self.store` for + // cross-batch evictions, or directly from `batch_entry` when + // a row evicts another row from the same in-flight batch + // (entry not yet registered in the store). + let evicted_batch = if prev_min.batch_id == batch_entry.id { + batch_entry.batch.clone() + } else { + self.store + .get(prev_min.batch_id) + .map(|entry| entry.batch.clone()) + .expect("evicted row's batch must be present in the store") + }; + let evicted = EvictedRow { + batch: evicted_batch, + index: prev_min.index, + row_bytes: prev_min.row.clone(), + }; + + // Update batch use + if prev_min.batch_id == batch_entry.id { + batch_entry.uses -= 1; + } else { + self.store.unuse(prev_min.batch_id); + } + + // update memory accounting + self.owned_bytes -= prev_min.owned_size(); + + prev_min.replace_with(row, batch_id, index); + + self.owned_bytes += prev_min.owned_size(); + + Some(evicted) + } else { + let new_row = TopKRow::new(row, batch_id, index); + self.owned_bytes += new_row.owned_size(); + // put the new row into the heap + self.inner.push(new_row); + None + } + } + + /// Returns the values stored in this heap, from values low to + /// high, as a single [`RecordBatch`], resetting the inner heap + pub fn emit(&mut self) -> Result<Option<RecordBatch>> { + Ok(self.emit_with_state()?.0) + } + + /// Returns the values stored in this heap, from values low to + /// high, as a single [`RecordBatch`], and a sorted vec of the + /// current heap's contents + fn emit_with_state(&mut self) -> Result<(Option<RecordBatch>, Vec<TopKRow>)> { + // generate sorted rows + let topk_rows = std::mem::take(&mut self.inner).into_sorted_vec(); + + if self.store.is_empty() { + return Ok((None, topk_rows)); + } + + // Collect the batches into a vec and store the "batch_id -> array_pos" mapping, to then + // build the `indices` vec below. This is needed since the batch ids are not continuous. + let mut record_batches = Vec::new(); + let mut batch_id_array_pos = HashMap::new(); + for (array_pos, (batch_id, batch)) in self.store.batches.iter().enumerate() { + record_batches.push(&batch.batch); + batch_id_array_pos.insert(*batch_id, array_pos); + } + + let indices: Vec<_> = topk_rows + .iter() + .map(|k| (batch_id_array_pos[&k.batch_id], k.index)) + .collect(); + + // At this point `indices` contains indexes within the + // rows and `input_arrays` contains a reference to the + // relevant RecordBatch for that index. `interleave_record_batch` pulls + // them together into a single new batch + let new_batch = interleave_record_batch(&record_batches, &indices)?; + + Ok((Some(new_batch), topk_rows)) + } + + /// Compact this heap, rewriting all stored batches into a single + /// input batch + pub fn maybe_compact(&mut self) -> Result<()> { + // Don't compact if there's only one batch (compacting into itself is pointless) + if self.store.len() <= 1 { + return Ok(()); + } + + let total_rows = self.store.total_rows; + let num_rows = self.inner.len(); + + // Compact when current store memory exceeds 2x what the compacted + // result would need. The multiplier avoids compacting when the + // savings would be marginal. + if total_rows <= num_rows * 2 { + return Ok(()); + } + + // at first, compact the entire thing always into a new batch + // (maybe we can get fancier in the future about ignoring + // batches that have a high usage ratio already + + // Note: new batch is in the same order as inner + let (new_batch, mut topk_rows) = self.emit_with_state()?; + let Some(new_batch) = new_batch else { + return Ok(()); + }; + + // clear all old entries in store (this invalidates all + // store_ids in `inner`) + self.store.clear(); + + let mut batch_entry = self.register_batch(new_batch); + batch_entry.uses = num_rows; + + // rewrite all existing entries to use the new batch, and + // remove old entries. The sortedness and their relative + // position do not change + for (i, topk_row) in topk_rows.iter_mut().enumerate() { + topk_row.batch_id = batch_entry.id; + topk_row.index = i; + } + self.insert_batch_entry(batch_entry); + // restore the heap + self.inner = BinaryHeap::from(topk_rows); + + Ok(()) + } + + /// return the size of memory used by this heap, in bytes + fn size(&self) -> usize { + size_of::<Self>() + + (self.inner.capacity() * size_of::<TopKRow>()) + + self.store.size() + + self.owned_bytes + } +} + +/// Represents one of the top K rows held in this heap. Orders +/// according to memcmp of row (e.g. the arrow Row format, but could +/// also be primitive values) +/// +/// Reuses allocations to minimize runtime overhead of creating new Vecs +#[derive(Debug, PartialEq)] +struct TopKRow { + /// the value of the sort key for this row. This contains the + /// bytes that could be stored in `OwnedRow` but uses `Vec<u8>` to + /// reuse allocations. + row: Vec<u8>, + /// the RecordBatch this row came from: an id into a [`RecordBatchStore`] + batch_id: u32, + /// the index in this record batch the row came from + index: usize, +} + +impl TopKRow { + /// Create a new TopKRow with new allocation + fn new(row: impl AsRef<[u8]>, batch_id: u32, index: usize) -> Self { + Self { + row: row.as_ref().to_vec(), + batch_id, + index, + } + } + + // Replace the existing row capacity with new values + fn replace_with(&mut self, new_row: impl AsRef<[u8]>, batch_id: u32, index: usize) { + self.row.clear(); + self.row.extend_from_slice(new_row.as_ref()); + + self.batch_id = batch_id; + self.index = index; + } + + /// Returns the number of bytes owned by this row in the heap (not + /// including itself) + fn owned_size(&self) -> usize { + self.row.capacity() + } + + /// Returns a slice to the owned row value + fn row(&self) -> &[u8] { + self.row.as_slice() + } +} + +impl Eq for TopKRow {} + +impl PartialOrd for TopKRow { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + // TODO PartialOrd is not consistent with PartialEq; PartialOrd contract is violated + Some(self.cmp(other)) + } +} + +impl Ord for TopKRow { + fn cmp(&self, other: &Self) -> Ordering { + self.row.cmp(&other.row) + } +} + +#[derive(Debug)] +struct RecordBatchEntry { + id: u32, + batch: RecordBatch, + // for this batch, how many times has it been used + uses: usize, +} + +/// This structure tracks [`RecordBatch`] by an id so that: +/// +/// 1. The baches can be tracked via an id that can be copied cheaply +/// 2. The total memory held by all batches is tracked +#[derive(Debug)] +struct RecordBatchStore { + /// id generator + next_id: u32, + /// storage + batches: HashMap<u32, RecordBatchEntry>, + /// total size of all record batches tracked by this store + batches_size: usize, + /// row count of all the batches + total_rows: usize, +} + +impl RecordBatchStore { + fn new() -> Self { + Self { + next_id: 0, + batches: HashMap::new(), + batches_size: 0, + total_rows: 0, + } + } + + /// Register this batch with the store and assign an ID. No + /// attempt is made to compare this batch to other batches + pub fn register(&mut self, batch: RecordBatch) -> RecordBatchEntry { + let id = self.next_id; + self.next_id += 1; + RecordBatchEntry { id, batch, uses: 0 } + } + + /// Insert a record batch entry into this store, tracking its + /// memory use, if it has any uses + pub fn insert(&mut self, entry: RecordBatchEntry) { + // uses of 0 means that none of the rows in the batch were stored in the topk + if entry.uses > 0 { + self.batches_size += get_record_batch_memory_size(&entry.batch); + self.total_rows += entry.batch.num_rows(); + self.batches.insert(entry.id, entry); + } + } + + /// Clear all values in this store, invalidating all previous batch ids + fn clear(&mut self) { + self.batches.clear(); + self.batches_size = 0; + self.total_rows = 0; + } + + fn get(&self, id: u32) -> Option<&RecordBatchEntry> { + self.batches.get(&id) + } + + /// returns the total number of batches stored in this store + fn len(&self) -> usize { + self.batches.len() + } + + /// returns true if the store has nothing stored + fn is_empty(&self) -> bool { + self.batches.is_empty() + } + + /// remove a use from the specified batch id. If the use count + /// reaches zero the batch entry is removed from the store + /// + /// panics if there were no remaining uses of id + pub fn unuse(&mut self, id: u32) { + let remove = if let Some(batch_entry) = self.batches.get_mut(&id) { + batch_entry.uses = batch_entry.uses.checked_sub(1).expect("underflow"); + batch_entry.uses == 0 + } else { + panic!("No entry for id {id}"); + }; + + if remove { + let old_entry = self.batches.remove(&id).unwrap(); + self.batches_size = self + .batches_size + .checked_sub(get_record_batch_memory_size(&old_entry.batch)) + .unwrap(); + + self.total_rows = self + .total_rows + .checked_sub(old_entry.batch.num_rows()) + .unwrap(); + } + } + + /// returns the size of memory used by this store, including all + /// referenced `RecordBatch`es, in bytes + pub fn size(&self) -> usize { + size_of::<Self>() + + self.batches.capacity() * (size_of::<u32>() + size_of::<RecordBatchEntry>()) + + self.batches_size + } +} + +/// Top-K-per-partition operator state. +/// +/// Sibling to [`TopK`]. Where `TopK` maintains a single global heap, +/// `PartitionedTopK` maintains one [`TopKHeap`] per distinct partition +/// key while sharing a single [`RowConverter`], [`MemoryReservation`], +/// scratch [`Rows`] buffer, and [`TopKMetrics`] across all partitions. +/// +/// This sharing is the point of the type: with N distinct partition +/// keys, a naive `HashMap<_, TopK>` pays N × constant overhead for +/// `RowConverter::new`, `MemoryConsumer::register`, and metric +/// counter setup. `PartitionedTopK` pays it once. +pub(crate) struct PartitionedTopK { + schema: SchemaRef, + metrics: TopKMetrics, + reservation: MemoryReservation, + /// ORDER BY expressions (excludes PARTITION BY). + expr: LexOrdering, + /// Encoder for ORDER BY columns. Reused across partitions. + row_converter: RowConverter, + /// Scratch row buffer reused across `insert_batch` calls. + scratch_rows: Rows, + /// PARTITION BY expressions. + partition_exprs: Vec<Arc<dyn PhysicalExpr>>, + /// Encoder for the partition key. + partition_converter: RowConverter, + /// One heap per distinct partition key seen so far. + heaps: HashMap<OwnedRow, TopKHeap>, + k: usize, + batch_size: usize, +} + +impl PartitionedTopK { + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_new( + partition_id: usize, + schema: SchemaRef, + partition_exprs: Vec<Arc<dyn PhysicalExpr>>, + partition_sort_fields: Vec<SortField>, + order_expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: &Arc<RuntimeEnv>, + metrics: &ExecutionPlanMetricsSet, + ) -> Result<Self> { + assert!(k > 0, "PartitionedTopK requires k > 0"); + let reservation = MemoryConsumer::new(format!("PartitionedTopK[{partition_id}]")) + .register(&runtime.memory_pool); + + let order_sort_fields = build_sort_fields(&order_expr, &schema)?; + let row_converter = RowConverter::new(order_sort_fields)?; + let scratch_rows = + row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + let partition_converter = RowConverter::new(partition_sort_fields)?; + + Ok(Self { + schema, + metrics: TopKMetrics::new(metrics, partition_id), + reservation, + expr: order_expr, + row_converter, + scratch_rows, + partition_exprs, + partition_converter, + heaps: HashMap::new(), + k, + batch_size, + }) + } + + /// Demultiplex `batch` rows by partition key, encode the ORDER BY + /// columns once for the whole batch, and feed each partition's + /// rows into its dedicated [`TopKHeap`]. + pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let baseline = self.metrics.baseline.clone(); + let _timer = baseline.elapsed_compute().timer(); + + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + + // 1. Evaluate + encode partition columns. + let pk_arrays: Vec<ArrayRef> = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::<Result<_>>()?; + let pk_rows = self.partition_converter.convert_columns(&pk_arrays)?; + + // 2. Demultiplex row indices by partition key (per-batch). + let mut groups: HashMap<OwnedRow, Vec<u32>> = HashMap::new(); + for i in 0..num_rows { + groups + .entry(pk_rows.row(i).owned()) + .or_default() + .push(i as u32); + } + + // 3. Evaluate ORDER BY columns on the full batch and encode ONCE. + let ob_arrays: Vec<ArrayRef> = self + .expr + .iter() + .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::<Result<_>>()?; + self.scratch_rows.clear(); + self.row_converter + .append(&mut self.scratch_rows, &ob_arrays)?; + + // 4. Per-partition: take the sub-batch, walk indices, dispatch + // qualifying rows into the partition's heap. + let k = self.k; + let mut replacements: usize = 0; + for (pk, indices) in groups { + let heap = self.heaps.entry(pk).or_insert_with(|| TopKHeap::new(k)); + + // Once a heap is full, most rows at high partition cardinality + // are rejected. Skip the gather + batch registration entirely + // when nothing in this partition group can improve the heap. + let any_qualify = indices.iter().any(|&orig_idx| { + let bytes = self.scratch_rows.row(orig_idx as usize); + match heap.max() { + Some(max_row) => bytes.as_ref() < max_row.row(), + None => true, + } + }); + if !any_qualify { + continue; + } + + let indices_arr = UInt32Array::from(indices); + let sub_batch = take_record_batch(batch, &indices_arr)?; + let mut entry = heap.register_batch(sub_batch); + + for (sub_idx, &orig_idx) in indices_arr.values().iter().enumerate() { + let row = self.scratch_rows.row(orig_idx as usize); + match heap.max() { + Some(max_row) if row.as_ref() >= max_row.row() => {} + None | Some(_) => { + heap.add(&mut entry, row, sub_idx); + replacements += 1; + } + } + } + + heap.insert_batch_entry(entry); + heap.maybe_compact()?; + } + + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + } + self.reservation.try_resize(self.size())?; + Ok(()) + } + + /// Drain all heaps in partition-key order and return the rows as + /// a stream of coalesced `RecordBatch`es ordered by + /// `(partition_keys, order_keys)`. + pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> { + let Self { + schema, + metrics, + reservation: _, + expr: _, + row_converter: _, + scratch_rows: _, + partition_exprs: _, + partition_converter: _, + mut heaps, + k: _, + batch_size, + } = self; + let _timer = metrics.baseline.elapsed_compute().timer(); + + let mut sorted_pks: Vec<OwnedRow> = heaps.keys().cloned().collect(); + sorted_pks.sort(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + + for pk in sorted_pks { + let mut heap = heaps.remove(&pk).expect("key from heaps.keys()"); + if let Some(batch) = heap.emit()? { + (&batch).record_output(&metrics.baseline); + coalescer.push_batch(batch)?; + } + } + coalescer.finish_buffered_batch()?; + + let mut out: Vec<Result<RecordBatch>> = Vec::new(); + while let Some(b) = coalescer.next_completed_batch() { + out.push(Ok(b)); + } + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(out), + ))) + } + + /// Total memory currently held by this operator, including all + /// per-partition heaps. + fn size(&self) -> usize { + size_of::<Self>() + + self.row_converter.size() + + self.partition_converter.size() + + self.scratch_rows.size() + + self.heaps.values().map(|h| h.size()).sum::<usize>() + + self.heaps.capacity() * (size_of::<OwnedRow>() + size_of::<TopKHeap>()) + } +} + +/// A run of rows from a single source [`RecordBatch`] that tied at the +/// boundary when inserted. Stored as `(batch, indices)` and materialized +/// at emit time via [`take_record_batch`]. +#[derive(Debug)] +struct TieEntry { + batch: RecordBatch, + /// Indices into `batch` of the rows tied at the (then-current) + /// boundary. Always non-empty by construction. + row_indices: Vec<u32>, + /// `get_record_batch_memory_size(&batch)` captured at push time so + /// `RankPartitionState::size()` doesn't recurse through `batch`'s + /// columns on every `try_resize` call. + batch_bytes: usize, +} + +/// Per-partition state for `RANK()` semantics. +/// +/// Composes [`TopKHeap`] as the K-bounded core plus a sibling +/// `Vec<TieEntry>` for boundary-tied rows. `RANK ≤ K` keeps the K +/// best rows by ORDER BY plus every row tied at the K-th-best +/// ORDER BY value — the boundary. So the total retained rows can +/// exceed K when ties straddle the boundary. +struct RankPartitionState { + heap: TopKHeap, + ties: Vec<TieEntry>, +} + +impl RankPartitionState { + fn size(&self) -> usize { + let ties_buffer = self.ties.capacity() * size_of::<TieEntry>(); + let ties_contents: usize = self + .ties + .iter() + .map(|t| t.row_indices.capacity() * size_of::<u32>() + t.batch_bytes) + .sum(); + self.heap.size() + ties_buffer + ties_contents + } +} + +/// Sibling to [`PartitionedTopK`] implementing `RANK()` semantics. +/// +/// Per partition, retains the K-best rows plus every row tied at the +/// K-th-best ORDER BY value (so `WHERE rk <= K` may keep more than K +/// rows when ties straddle the boundary). Like [`PartitionedTopK`], +/// the [`RowConverter`], [`MemoryReservation`], scratch [`Rows`] +/// buffer, and [`TopKMetrics`] are shared across all partitions for +/// this operator instance. +/// +/// # Algorithm (per row) +/// +/// For each incoming row, compare its encoded ORDER BY bytes against +/// `heap.max()` — the K-th-best row, which is by definition the +/// admission boundary. `heap.max()` is `None` until the heap fills +/// to K rows: +/// +/// - heap not full (`max() == None`) → forward to the heap +/// - row's ob `==` max → push to ties (no heap call) +/// - row's ob `>` max → drop +/// - row's ob `<` max → forward to heap; on eviction, compare the +/// new `heap.max()` to the evicted row's bytes: if equal, push +/// evicted to ties (still tied at the new boundary's rank); else +/// clear ties (boundary moved up, old ties no longer satisfy +/// `rk ≤ K`) +pub(crate) struct PartitionedTopKRank { + schema: SchemaRef, + metrics: TopKMetrics, + reservation: MemoryReservation, + /// ORDER BY expressions (excludes PARTITION BY). + expr: LexOrdering, + /// Encoder for ORDER BY columns. Reused across partitions. + row_converter: RowConverter, + /// Scratch row buffer reused across `insert_batch` calls. + scratch_rows: Rows, + /// PARTITION BY expressions. + partition_exprs: Vec<Arc<dyn PhysicalExpr>>, + /// Encoder for the partition key. + partition_converter: RowConverter, + /// Scratch row buffer for partition-key encoding. Reused across + /// `insert_batch` calls (cleared + appended each batch) so we + /// avoid allocating a fresh `Rows` buffer every batch. + partition_scratch_rows: Rows, + /// One rank state per distinct partition key seen so far. + states: HashMap<OwnedRow, RankPartitionState>, + k: usize, + batch_size: usize, +} + +impl PartitionedTopKRank { + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_new( + partition_id: usize, + schema: SchemaRef, + partition_exprs: Vec<Arc<dyn PhysicalExpr>>, + partition_sort_fields: Vec<SortField>, + order_expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: &Arc<RuntimeEnv>, + metrics: &ExecutionPlanMetricsSet, + ) -> Result<Self> { + assert!(k > 0, "PartitionedTopKRank requires k > 0"); + let reservation = + MemoryConsumer::new(format!("PartitionedTopKRank[{partition_id}]")) + .register(&runtime.memory_pool); + + let order_sort_fields = build_sort_fields(&order_expr, &schema)?; + let row_converter = RowConverter::new(order_sort_fields)?; + let scratch_rows = + row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + let partition_converter = RowConverter::new(partition_sort_fields)?; + let partition_scratch_rows = partition_converter + .empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + Ok(Self { + schema, + metrics: TopKMetrics::new(metrics, partition_id), + reservation, + expr: order_expr, + row_converter, + scratch_rows, + partition_exprs, + partition_converter, + partition_scratch_rows, + states: HashMap::new(), + k, + batch_size, + }) + } + + /// Demultiplex `batch` rows by partition key, encode the ORDER BY + /// columns once for the whole batch, and feed each partition's + /// rows through the rank classifier into its dedicated heap and + /// ties Vec. + pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let baseline = self.metrics.baseline.clone(); + let _timer = baseline.elapsed_compute().timer(); + + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + + // Captured once so the per-tie push from this batch can reuse + // it (computing `get_record_batch_memory_size` is O(cols × + // buffer walk) and we'd otherwise pay it per push and again + // per `try_resize` call). + let input_batch_bytes = get_record_batch_memory_size(batch); + + // 1. Evaluate + encode partition columns into the reusable + // scratch (cleared then appended). + let pk_arrays: Vec<ArrayRef> = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::<Result<_>>()?; + self.partition_scratch_rows.clear(); + self.partition_converter + .append(&mut self.partition_scratch_rows, &pk_arrays)?; + let pk_rows = &self.partition_scratch_rows; + + // 2. Demultiplex row indices by partition key (per-batch). + let mut groups: HashMap<OwnedRow, Vec<u32>> = HashMap::new(); + for i in 0..num_rows { + groups + .entry(pk_rows.row(i).owned()) + .or_default() + .push(i as u32); + } + + // 3. Evaluate ORDER BY columns on the full batch and encode ONCE. + let ob_arrays: Vec<ArrayRef> = self + .expr + .iter() + .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::<Result<_>>()?; + self.scratch_rows.clear(); + self.row_converter + .append(&mut self.scratch_rows, &ob_arrays)?; + + // 4. Per-partition: classify each row and dispatch. + let k = self.k; + let mut replacements: usize = 0; + + for (pk, indices) in groups { + let state = self.states.entry(pk).or_insert_with(|| RankPartitionState { + heap: TopKHeap::new(k), + ties: Vec::new(), + }); + + // Equal indices for THIS batch only. Coalesced into a single + // tie entry at the end of the partition's loop. Discarded if + // the boundary moves up mid-loop (those rows were tied to the + // old boundary, which is now strictly worse than the new K-th). + let mut equal_indices: Vec<u32> = Vec::new(); + // Lazy-registered: only attached if at least one row reaches + // the heap from this batch in this partition. + let mut entry: Option<RecordBatchEntry> = None; + + for &orig_idx in &indices { + let row = self.scratch_rows.row(orig_idx as usize); + + // Classify against the current K-th-best (the heap top). + // `heap.max()` returns `None` while the heap is filling, + // so unclassified rows fall through to the heap path. + let classification = state + .heap + .max() + .map(|max_row| row.as_ref().cmp(max_row.row())); + + match classification { + Some(Ordering::Equal) => { + equal_indices.push(orig_idx); + continue; + } + Some(Ordering::Greater) => continue, + Some(Ordering::Less) | None => { + // Heap path: heap not yet full, or row strictly + // better than the current boundary. + let entry_ref = entry.get_or_insert_with(|| { + state.heap.register_batch(batch.clone()) + }); + if let Some(EvictedRow { + batch: evicted_batch, + index: evicted_index, + row_bytes: evicted_bytes, + }) = state.heap.add(entry_ref, row, orig_idx as usize) + { + // Compare the new boundary (post-eviction heap + // top) against the evicted row's bytes — both + // already in encoded form, no clones needed. + let boundary_changed = state + .heap + .max() + .expect("heap was full to evict; must still be full") + .row() + != evicted_bytes.as_slice(); + if boundary_changed { + // Boundary moved up — prior ties (across + // all prior batches) and equal_indices + // accumulated earlier in THIS batch were + // tied to the old boundary, now strictly + // worse than the new K-th-best. Discard. + state.ties.clear(); + equal_indices.clear(); + } else { + // Boundary unchanged — evicted row is tied + // at the (unchanged) boundary; push as a + // single-row entry. + let batch_bytes = + get_record_batch_memory_size(&evicted_batch); + state.ties.push(TieEntry { + batch: evicted_batch, + row_indices: vec![evicted_index as u32], + batch_bytes, + }); + } + } + replacements += 1; + } + } + } + + if let Some(e) = entry { + state.heap.insert_batch_entry(e); + state.heap.maybe_compact()?; + } + + // Commit this batch's ties as a single entry. + if !equal_indices.is_empty() { + state.ties.push(TieEntry { + batch: batch.clone(), + row_indices: equal_indices, + batch_bytes: input_batch_bytes, + }); + } + } + + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + } + self.reservation.try_resize(self.size())?; + Ok(()) + } + + /// Drain all heaps and ties in partition-key order and return the + /// rows as a stream of coalesced [`RecordBatch`]es ordered by + /// `(partition_keys, order_keys)`. Within a partition, heap rows + /// come first (sorted by ob), then tie rows (all sharing the + /// boundary ob). + pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> { + let Self { + schema, + metrics, + reservation: _, + expr: _, + row_converter: _, + scratch_rows: _, + partition_exprs: _, + partition_converter: _, + partition_scratch_rows: _, + mut states, + k: _, + batch_size, + } = self; + let _timer = metrics.baseline.elapsed_compute().timer(); + + let mut sorted_pks: Vec<OwnedRow> = states.keys().cloned().collect(); + sorted_pks.sort(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + + for pk in sorted_pks { + let RankPartitionState { mut heap, ties, .. } = + states.remove(&pk).expect("key from states.keys()"); + if let Some(batch) = heap.emit()? { + (&batch).record_output(&metrics.baseline); + coalescer.push_batch(batch)?; + } + for tie in ties { + let indices = UInt32Array::from(tie.row_indices); + let tie_batch = take_record_batch(&tie.batch, &indices)?; + (&tie_batch).record_output(&metrics.baseline); + coalescer.push_batch(tie_batch)?; + } + } + coalescer.finish_buffered_batch()?; + + let mut out: Vec<Result<RecordBatch>> = Vec::new(); + while let Some(b) = coalescer.next_completed_batch() { + out.push(Ok(b)); + } + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(out), + ))) + } + + /// Total memory currently held, including all per-partition states. + fn size(&self) -> usize { + size_of::<Self>() + + self.row_converter.size() + + self.partition_converter.size() + + self.scratch_rows.size() + + self.partition_scratch_rows.size() + + self.states.values().map(|s| s.size()).sum::<usize>() + + self.states.capacity() + * (size_of::<OwnedRow>() + size_of::<RankPartitionState>()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{BooleanArray, Float64Array, Int32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow_schema::SortOptions; + use datafusion_common::assert_batches_eq; + use datafusion_physical_expr::{DynamicFilterTracking, expressions::col}; + use futures::TryStreamExt; + + /// This test ensures the size calculation is correct for RecordBatches with multiple columns. + #[test] + fn test_record_batch_store_size() { + // given + let schema = Arc::new(Schema::new(vec![ + Field::new("ints", DataType::Int32, true), + Field::new("float64", DataType::Float64, false), + ])); + let mut record_batch_store = RecordBatchStore::new(); + let int_array = + Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); // 5 * 4 = 20 + let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]); // 5 * 8 = 40 + + let record_batch_entry = RecordBatchEntry { + id: 0, + batch: RecordBatch::try_new( + schema, + vec![Arc::new(int_array), Arc::new(float64_array)], + ) + .unwrap(), + uses: 1, + }; + + // when insert record batch entry + record_batch_store.insert(record_batch_entry); + assert_eq!(record_batch_store.batches_size, 60); + + // when unuse record batch entry + record_batch_store.unuse(0); + assert_eq!(record_batch_store.batches_size, 0); + } + + fn make_ab_schema() -> SchemaRef { + make_ab_schema_with_nullable_a(false) + } + + fn make_ab_schema_with_nullable_a(a_nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, a_nullable), + Field::new("b", DataType::Float64, false), + ])) + } + + // Local TopK tests use one emitter; shared-filter cases pass the partition count explicitly. + fn make_topk_filter() -> Arc<RwLock<TopKDynamicFilters>> { + make_shared_topk_filter(1) + } + + fn make_shared_topk_filter( + topk_emitter_count: usize, + ) -> Arc<RwLock<TopKDynamicFilters>> { + Arc::new(RwLock::new( + TopKDynamicFilters::new_with_topk_emitter_count( + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))), + topk_emitter_count, + ), + )) + } + + /// Builds the `(a, b)` fixture used by prefix-completion tests: + /// full sort `(a, b)`, input prefix `[a]`, `k = 3`, and batch size 2. + fn make_ab_topk( + schema: SchemaRef, + filter: Arc<RwLock<TopKDynamicFilters>>, + ) -> Result<TopK> { + make_ab_topk_with_options(0, schema, filter, SortOptions::default()) + } + + fn make_ab_topk_with_options( + partition_id: usize, + schema: SchemaRef, + filter: Arc<RwLock<TopKDynamicFilters>>, + a_options: SortOptions, + ) -> Result<TopK> { + let sort_expr_a = PhysicalSortExpr { + expr: col("a", schema.as_ref())?, + options: a_options, + }; + let sort_expr_b = PhysicalSortExpr { + expr: col("b", schema.as_ref())?, + options: SortOptions::default(), + }; + + TopK::try_new( + partition_id, + schema, + vec![sort_expr_a.clone()], + LexOrdering::from([sort_expr_a, sort_expr_b]), + 3, + 2, + Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + filter, + ) + } + + fn make_ab_batch( + schema: SchemaRef, + a: &[Option<i32>], + b: &[f64], + ) -> Result<RecordBatch> { + Ok(RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(a.to_vec())) as ArrayRef, + Arc::new(Float64Array::from(b.to_vec())) as ArrayRef, + ], + )?) + } + + type AbRow = (Option<i32>, f64); + + fn make_ab_rows_batch(schema: SchemaRef, rows: &[AbRow]) -> Result<RecordBatch> { + let (a, b): (Vec<_>, Vec<_>) = rows.iter().copied().unzip(); + make_ab_batch(schema, &a, &b) + } + + #[tokio::test] + async fn test_early_completion_marks_finished_with_prefix() -> Result<()> { + let schema = make_ab_schema(); + let mut topk = make_ab_topk(Arc::clone(&schema), make_topk_filter())?; + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(1), Some(1), Some(2)], + &[20.0, 15.0, 30.0], + )?)?; + assert!(!topk.finished); + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(2), Some(3)], + &[10.0, 20.0], + )?)?; + assert!(topk.finished); + + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+---+------+", + "| a | b |", + "+---+------+", + "| 1 | 15.0 |", + "| 1 | 20.0 |", + "| 2 | 10.0 |", + "+---+------+", + ], + &results + ); + + Ok(()) + } + + /// Regression test for #22849: a batch whose rows are entirely rejected by the + /// heap's dynamic filter must still trigger `attempt_early_completion` when its + /// last row's prefix is worse than the heap's worst. + #[tokio::test] + async fn test_early_completion_fires_when_filter_rejects_entire_batch() -> Result<()> + { + let schema = make_ab_schema(); + let mut topk = make_ab_topk(Arc::clone(&schema), make_topk_filter())?; + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(1), Some(1), Some(2)], + &[20.0, 15.0, 30.0], + )?)?; + assert!(!topk.finished); + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(3), Some(3)], + &[10.0, 20.0], + )?)?; + assert!(topk.finished); + + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+---+------+", + "| a | b |", + "+---+------+", + "| 1 | 15.0 |", + "| 1 | 20.0 |", + "| 2 | 30.0 |", + "+---+------+", + ], + &results + ); + + Ok(()) + } + + #[tokio::test] + async fn test_early_completion_fires_when_batch_makes_no_replacements() -> Result<()> + { + let schema = make_ab_schema(); + let filter = make_topk_filter(); + let mut topk = make_ab_topk(Arc::clone(&schema), Arc::clone(&filter))?; + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(1), Some(1), Some(2)], + &[20.0, 15.0, 30.0], + )?)?; + assert!(!topk.finished); + + let replacements_before = topk.metrics.row_replacements.value(); + + // Keep the dynamic filter permissive so the second batch reaches + // `find_new_topk_items`; all of its rows are worse than the heap max, + // so this specifically exercises the `replacements == 0` path. + filter.read().expr().update(lit(true))?; + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(3), Some(3)], + &[10.0, 20.0], + )?)?; + assert_eq!(topk.metrics.row_replacements.value(), replacements_before); + assert!(topk.finished); + + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+---+------+", + "| a | b |", + "+---+------+", + "| 1 | 15.0 |", + "| 1 | 20.0 |", + "| 2 | 30.0 |", + "+---+------+", + ], + &results + ); + + Ok(()) + } + + struct SharedPrefixCase { + name: &'static str, + a_nullable: bool, + a_options: SortOptions, + threshold_source_rows: &'static [AbRow], + lagging_partition_rows: &'static [AbRow], + expected_finished: bool, + } + + fn assert_shared_prefix_case(case: SharedPrefixCase) -> Result<()> { + let schema = make_ab_schema_with_nullable_a(case.a_nullable); + let filter = make_shared_topk_filter(2); + + let mut threshold_source = make_ab_topk_with_options( + 0, + Arc::clone(&schema), + Arc::clone(&filter), + case.a_options, + )?; + threshold_source.insert_batch(make_ab_rows_batch( + Arc::clone(&schema), + case.threshold_source_rows, + )?)?; + assert!( + filter + .read() + .shared_threshold + .as_ref() + .and_then(TopKThreshold::common_prefix_row) + .is_some(), + "{}: threshold-source partition should establish the shared prefix threshold", + case.name + ); + + let mut lagging_partition = make_ab_topk_with_options( + 1, + Arc::clone(&schema), + Arc::clone(&filter), + case.a_options, + )?; + lagging_partition + .insert_batch(make_ab_rows_batch(schema, case.lagging_partition_rows)?)?; + + assert!( + lagging_partition.heap.inner.is_empty(), + "{}: lagging partition's local heap should remain empty", + case.name + ); + assert_eq!( + lagging_partition.finished, case.expected_finished, + "{}", + case.name + ); + + Ok(()) + } + + #[test] + fn test_shared_filter_can_finish_partition_before_local_heap_is_full() -> Result<()> { + assert_shared_prefix_case(SharedPrefixCase { + name: "shared threshold should finish lagging partition", + a_nullable: false, + a_options: SortOptions::default(), + threshold_source_rows: &[(Some(1), 20.0), (Some(1), 15.0), (Some(2), 30.0)], + lagging_partition_rows: &[(Some(3), 10.0), (Some(3), 20.0)], + expected_finished: true, + }) + } + + #[test] + fn test_shared_prefix_threshold_boundary_cases() -> Result<()> { + for case in [ + SharedPrefixCase { + name: "equal prefix cannot prove completion", + a_nullable: false, + a_options: SortOptions::default(), + threshold_source_rows: &[ + (Some(1), 20.0), + (Some(1), 15.0), + (Some(2), 30.0), + ], + lagging_partition_rows: &[(Some(2), 40.0), (Some(2), 50.0)], + expected_finished: false, + }, + SharedPrefixCase { + name: "descending prefix uses sort-order row encoding", + a_nullable: false, + a_options: SortOptions { + descending: true, + nulls_first: true, + }, + threshold_source_rows: &[ + (Some(10), 1.0), + (Some(10), 2.0), + (Some(9), 3.0), + ], + lagging_partition_rows: &[(Some(8), 1.0), (Some(8), 2.0)], + expected_finished: true, + }, + SharedPrefixCase { + name: "NULLS LAST prefix uses sort-order row encoding", + a_nullable: true, + a_options: SortOptions { + descending: false, + nulls_first: false, + }, + threshold_source_rows: &[ + (Some(1), 20.0), + (Some(1), 15.0), + (Some(2), 30.0), + ], + lagging_partition_rows: &[(None, 10.0), (None, 20.0)], + expected_finished: true, + }, + ] { + assert_shared_prefix_case(case)?; + } + Ok(()) + } + + fn make_single_column_topk( + dynamic_filter: Arc<DynamicFilterPhysicalExpr>, + ) -> Result<(SchemaRef, TopK)> { + make_single_column_topk_with_filter( + 0, + Arc::new(RwLock::new(TopKDynamicFilters::new(dynamic_filter))), + ) + } + + fn make_single_column_topk_with_filter( + partition_id: usize, + filter: Arc<RwLock<TopKDynamicFilters>>, + ) -> Result<(SchemaRef, TopK)> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let sort_expr = PhysicalSortExpr { + expr: col("a", schema.as_ref())?, + options: SortOptions::default(), + }; + + let topk = TopK::try_new( + partition_id, + Arc::clone(&schema), + vec![sort_expr.clone()], + LexOrdering::from([sort_expr]), + 2, + 10, + Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + filter, + )?; + + Ok((schema, topk)) + } + + #[tokio::test] + async fn test_topk_marks_filter_complete() -> Result<()> { + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let dynamic_filter_clone = Arc::clone(&dynamic_filter); + let (schema, mut topk) = make_single_column_topk(dynamic_filter)?; + + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), Some(1), Some(2)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array])?; + topk.insert_batch(batch)?; + + let _results: Vec<_> = topk.emit()?.try_collect().await?; + + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter_clone.wait_complete(), + ) + .await + .expect("single-emitter TopK should mark the dynamic filter complete"); + + Ok(()) + } + + #[tokio::test] + async fn test_shared_topk_filter_completes_after_last_emitter() -> Result<()> { + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let dynamic_filter_clone = Arc::clone(&dynamic_filter); + let shared_filter = Arc::new(RwLock::new( + TopKDynamicFilters::new_with_topk_emitter_count(dynamic_filter, 2), + )); + + let (schema, mut topk_0) = + make_single_column_topk_with_filter(0, Arc::clone(&shared_filter))?; + let (_, mut topk_1) = + make_single_column_topk_with_filter(1, Arc::clone(&shared_filter))?; + + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), Some(1), Some(2)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array])?; + topk_0.insert_batch(batch)?; + let _results: Vec<_> = topk_0.emit()?.try_collect().await?; + + let dynamic_filter_expr: Arc<dyn PhysicalExpr> = + Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter_clone); + assert!( + matches!( + DynamicFilterTracking::classify(&dynamic_filter_expr), + DynamicFilterTracking::Watching(_) + ), + "the shared filter should remain watchable until every TopK emits" + ); + + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(6), Some(4), Some(5)])); + let batch = RecordBatch::try_new(schema, vec![array])?; + topk_1.insert_batch(batch)?; + let _results: Vec<_> = topk_1.emit()?.try_collect().await?; + + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter_clone.wait_complete(), + ) + .await + .expect("the final shared TopK emitter should mark the dynamic filter complete"); + + Ok(()) + } + + /// Tests that memory-based compaction triggers when a large batch + /// has very few rows referenced by the top-k heap. + #[tokio::test] + async fn test_topk_memory_compaction() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + + let sort_expr = PhysicalSortExpr { + expr: col("a", schema.as_ref())?, + options: SortOptions::default(), + }; + + let full_expr = LexOrdering::from([sort_expr.clone()]); + let prefix = vec![sort_expr]; + + let runtime = Arc::new(RuntimeEnv::default()); + let metrics = ExecutionPlanMetricsSet::new(); + + let k = 5; + let mut topk = TopK::try_new( + 0, + Arc::clone(&schema), + prefix, + full_expr, + k, + 8192, + runtime, + &metrics, + Arc::new(RwLock::new(TopKDynamicFilters::new(Arc::new( + DynamicFilterPhysicalExpr::new(vec![], lit(true)), + )))), + )?; + + // Insert a large batch (100,000 rows) with values 1..=100_000. + // Only the smallest 5 values (1..=5) will end up in the heap. + let large_values: Vec<i32> = (1..=100_000).collect(); + let array1: ArrayRef = Arc::new(Int32Array::from(large_values)); + let batch1 = RecordBatch::try_new(Arc::clone(&schema), vec![array1])?; + topk.insert_batch(batch1)?; + + // After the first batch, store has 1 batch — compaction should + // not trigger (guard: store.len() <= 1). + assert_eq!( + topk.heap.store.len(), + 1, + "should have 1 batch before second insert" + ); + + // Insert a second batch whose values displace entries in the heap. + // -1 and 0 are smaller than the current top-5 (1..=5), so they + // produce 2 replacements. With replacements > 0, `insert_batch` + // calls `insert_batch_entry` (briefly making store.len() == 2) + // and then `maybe_compact`, which should collapse it back to 1. + let array2: ArrayRef = Arc::new(Int32Array::from(vec![-1, 0])); + let batch2 = RecordBatch::try_new(Arc::clone(&schema), vec![array2])?; + let replacements_before = topk.metrics.row_replacements.value(); + topk.insert_batch(batch2)?; + + // Sanity check: batch2 was actually integrated. Without + // replacements, `maybe_compact` is never called and the + // store-length assertion below would pass vacuously. + assert!( + topk.metrics.row_replacements.value() > replacements_before, + "batch2 must produce replacements so compaction is exercised" + ); + + // The compacted-estimate guard is `total_rows <= num_rows * 2`, + // i.e. 100_002 <= 10, which is false, so compaction fires and + // collapses the two stored batches back into one. + assert_eq!( + topk.heap.store.len(), + 1, + "store should be compacted to 1 batch" + ); + + // Verify the emitted results are correct (top 5 ascending). + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+", "| a |", "+----+", "| -1 |", "| 0 |", "| 1 |", "| 2 |", + "| 3 |", "+----+", + ], + &results + ); + + Ok(()) + } + + /// Negative path: when stored rows are close to the heap size, + /// compaction must NOT fire even with multiple batches present, + /// because the savings would be marginal + /// (guard: `total_rows <= num_rows * 2`). + /// + /// Uses a bit-packed `BooleanArray` so that future changes to the + /// compaction heuristic that reintroduce a per-byte estimate + /// (where integer truncation could misbehave on sub-byte types) + /// are caught here. + #[tokio::test] + async fn test_topk_memory_compaction_skipped_when_marginal() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Boolean, false)])); + + let sort_expr = PhysicalSortExpr { + expr: col("a", schema.as_ref())?, + options: SortOptions::default(), + }; + let full_expr = LexOrdering::from([sort_expr.clone()]); + let prefix = vec![sort_expr]; + + let runtime = Arc::new(RuntimeEnv::default()); + let metrics = ExecutionPlanMetricsSet::new(); + + let k = 10; + let mut topk = TopK::try_new( + 0, + Arc::clone(&schema), + prefix, + full_expr, + k, + 8192, + runtime, + &metrics, + Arc::new(RwLock::new(TopKDynamicFilters::new(Arc::new( + DynamicFilterPhysicalExpr::new(vec![], lit(true)), + )))), + )?; + + // Two small batches; every row from both batches ends up referenced + // by the heap, so total_rows == num_rows == 10. + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(BooleanArray::from(vec![false, false, true, true, true])) + as ArrayRef, + ], + )?; + topk.insert_batch(batch1)?; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(BooleanArray::from(vec![false, false, false, true, true])) + as ArrayRef, + ], + )?; + topk.insert_batch(batch2)?; + + // Guard `total_rows <= num_rows * 2` should hold (10 <= 20), + // so compaction is skipped and BOTH batches remain in the store. + assert_eq!( + topk.heap.store.len(), + 2, + "store must keep 2 batches when savings would be marginal" + ); + assert_eq!(topk.heap.inner.len(), 10, "heap should hold all 10 rows"); + + // Output is still correct (5 falses then 5 trues ascending). + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+-------+", + "| a |", + "+-------+", + "| false |", + "| false |", + "| false |", + "| false |", + "| false |", + "| true |", + "| true |", + "| true |", + "| true |", + "| true |", + "+-------+", + ], + &results + ); + + Ok(()) + } + + /// Builds a `(pk Int32, val Int32)` schema and a `PartitionedTopK` + /// partitioned by `pk` with order `val ASC`. Helper for the + /// `PartitionedTopK` tests below. + fn build_partitioned_topk(k: usize) -> Result<(Arc<Schema>, PartitionedTopK)> { + build_partitioned_topk_with_opts(k, SortOptions::default(), false) + } + + /// Variant of [`build_partitioned_topk`] that lets the test pick the + /// `val` column's `SortOptions` (direction, null ordering) and + /// nullability. Used by tests that exercise the shared encoder + /// across `ASC`/`DESC` and `NULLS FIRST/LAST` paths. + fn build_partitioned_topk_with_opts( + k: usize, + val_sort_options: SortOptions, + val_nullable: bool, + ) -> Result<(Arc<Schema>, PartitionedTopK)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Int32, val_nullable), + ])); + + let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?; + let pk_sort_expr = PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }; + let val_sort_expr = PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: val_sort_options, + }; + + let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?; + let order_expr = LexOrdering::from([val_sort_expr]); + + let state = PartitionedTopK::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + k, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + Ok((schema, state)) + } + + fn pk_val_batch( + schema: &Arc<Schema>, + pks: Vec<i32>, + vals: Vec<i32>, + ) -> Result<RecordBatch> { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(pks)), + Arc::new(Int32Array::from(vals)), + ], + )?) + } + + /// Variant of [`pk_val_batch`] that accepts nullable `val`s. Used by + /// tests that exercise null-ordering through the shared encoder. + fn nullable_pk_val_batch( + schema: &Arc<Schema>, + pks: Vec<i32>, + vals: Vec<Option<i32>>, + ) -> Result<RecordBatch> { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(pks)), + Arc::new(Int32Array::from(vals)), + ], + )?) + } + + /// Multiple distinct partition keys interleaved within a single + /// input batch — the per-batch demux, per-partition heap eviction, + /// and partition-key-ordered emit must all behave correctly. + #[tokio::test] + async fn test_partitioned_topk_multi_partition_within_batch() -> Result<()> { + let (schema, mut state) = build_partitioned_topk(2)?; + + // pk=1 vals: 10, 5, 8 → top-2 ASC = [5, 8] + // pk=2 vals: 20, 15 → top-2 ASC = [15, 20] + // pk=3 vals: 7 → top-2 ASC = [7] + let batch = + pk_val_batch(&schema, vec![1, 2, 1, 2, 1, 3], vec![10, 20, 5, 15, 8, 7])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 8 |", + "| 2 | 15 |", + "| 2 | 20 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// State must accumulate across `insert_batch` calls: a partition + /// key seen in batch 1 should still own its heap when batch 2 + /// arrives, and a row in batch 2 that beats the existing K-th + /// best should evict the loser. + #[tokio::test] + async fn test_partitioned_topk_cross_batch_eviction() -> Result<()> { + let (schema, mut state) = build_partitioned_topk(2)?; + + // Batch 1: pk=1 fills the heap with [50, 40]. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 40])?)?; + + // Batch 2: pk=1 sees a smaller value (10) — it must evict 50. + // pk=2 appears for the first time mid-stream. + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 2, 1], + vec![10, 99, 60], // 60 > 40 stays on top, gets discarded + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 10 |", + "| 1 | 40 |", + "| 2 | 99 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Empty input must produce an empty output stream, not panic. + #[tokio::test] + async fn test_partitioned_topk_empty_input() -> Result<()> { + let (_schema, state) = build_partitioned_topk(3)?; + let results: Vec<_> = state.emit()?.try_collect().await?; + assert!(results.is_empty(), "empty input → empty output"); + Ok(()) + } + + /// `fetch = 1` is a common case (rn = 1 filter). The heap should + /// hold exactly one row per partition: the partition's minimum. + #[tokio::test] + async fn test_partitioned_topk_fetch_one() -> Result<()> { + let (schema, mut state) = build_partitioned_topk(1)?; + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 1, 2, 2, 3], + vec![3, 1, 9, 4, 7], + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 1 |", + "| 2 | 4 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ORDER BY val DESC` exercises the shared encoder's sort-direction + /// handling: the row converter must flip the sort sign for `val` so + /// that larger values compare smaller in row-encoded form. Each + /// partition should keep its top-K *largest* values. + #[tokio::test] + async fn test_partitioned_topk_desc_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_with_opts( + 2, + SortOptions { + descending: true, + nulls_first: false, + }, + false, + )?; + + // pk=1 vals: 10, 5, 8, 12 → top-2 DESC = [12, 10] + // pk=2 vals: 20, 15, 25 → top-2 DESC = [25, 20] + let batch = pk_val_batch( + &schema, + vec![1, 2, 1, 2, 1, 1, 2], + vec![10, 20, 5, 15, 8, 12, 25], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 12 |", + "| 1 | 10 |", + "| 2 | 25 |", + "| 2 | 20 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// NULL sort values exercise the shared encoder's null-ordering + /// handling. With `ASC NULLS LAST`, NULLs sort *after* every + /// non-NULL value, so a partition whose only non-NULL value beats + /// a NULL must evict the NULL when `K = 1`. A partition that holds + /// only NULLs must still emit them. + #[tokio::test] + async fn test_partitioned_topk_nulls_last_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_with_opts( + 1, + SortOptions { + descending: false, + nulls_first: false, + }, + true, + )?; + + // pk=1 vals: NULL, 7, NULL → top-1 ASC NULLS LAST = [7] + // pk=2 vals: NULL → top-1 = [NULL] + // pk=3 vals: NULL, 4, 2 → top-1 = [2] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 1, 3, 3, 3], + vec![None, None, Some(7), None, None, Some(4), Some(2)], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 7 |", + "| 2 | |", + "| 3 | 2 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ASC NULLS FIRST` (the `SortOptions::default()`) sorts NULLs + /// *before* every non-NULL value, so under `fetch = K` a partition's + /// NULLs are kept preferentially over larger non-NULL values. + #[tokio::test] + async fn test_partitioned_topk_nulls_first_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_with_opts( + 2, + SortOptions { + descending: false, + nulls_first: true, + }, + true, + )?; + + // pk=1 vals: NULL, 5, NULL, 8 → top-2 ASC NULLS FIRST = [NULL, NULL] + // pk=2 vals: 7, NULL → top-2 = [NULL, 7] + // pk=3 vals: 3, 1 → top-2 = [1, 3] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 3, 1, 2, 1, 3], + vec![ + None, + Some(7), + Some(5), + Some(3), + None, + None, + Some(8), + Some(1), + ], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | |", + "| 1 | |", + "| 2 | |", + "| 2 | 7 |", + "| 3 | 1 |", + "| 3 | 3 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + // ==================================================================== + // PartitionedTopKRank operator tests + // + // These mirror the PartitionedTopK tests above plus three RANK-specific + // cases for the Equal / boundary-shift / boundary-unchanged-eviction + // arms in `PartitionedTopKRank::insert_batch`. + // ==================================================================== + + /// Builds a `(pk Int32, val Int32)` schema and a `PartitionedTopKRank` + /// keyed on `pk ASC` (partition) and `val ASC` (ORDER BY). + fn build_partitioned_topk_rank( + k: usize, + ) -> Result<(Arc<Schema>, PartitionedTopKRank)> { + build_partitioned_topk_rank_with_opts(k, SortOptions::default(), false) + } + + /// Variant of [`build_partitioned_topk_rank`] that lets the test pick + /// the `val` column's `SortOptions` (direction, null ordering) and + /// nullability. + fn build_partitioned_topk_rank_with_opts( + k: usize, + val_sort_options: SortOptions, + val_nullable: bool, + ) -> Result<(Arc<Schema>, PartitionedTopKRank)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Int32, val_nullable), + ])); + + let pk_expr: Arc<dyn PhysicalExpr> = col("pk", schema.as_ref())?; + let pk_sort_expr = PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }; + let val_sort_expr = PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: val_sort_options, + }; + + let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?; + let order_expr = LexOrdering::from([val_sort_expr]); + + let state = PartitionedTopKRank::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + k, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + Ok((schema, state)) + } + + /// Multiple distinct partition keys interleaved within a single + /// input batch — the per-batch demux, per-partition heap eviction, + /// and partition-key-ordered emit must all behave correctly. No + /// ties: result should match a `ROW_NUMBER` top-K under the same K. + #[tokio::test] + async fn test_partitioned_topk_rank_multi_partition_within_batch() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 5, 8 → top-2 ASC = [5, 8] + // pk=2 vals: 20, 15 → top-2 ASC = [15, 20] + // pk=3 vals: 7 → top-2 ASC = [7] + let batch = + pk_val_batch(&schema, vec![1, 2, 1, 2, 1, 3], vec![10, 20, 5, 15, 8, 7])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 8 |", + "| 2 | 15 |", + "| 2 | 20 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// State must accumulate across `insert_batch` calls. A row in + /// batch 2 that's strictly better than the existing K-th must + /// evict it; an evicted row whose bytes match the new boundary + /// becomes a `TieEntry` pinned to the prior batch. + #[tokio::test] + async fn test_partitioned_topk_rank_cross_batch_eviction() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // Batch 1: pk=1 fills the heap with [50, 40]. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 40])?)?; + + // Batch 2: pk=1 sees a smaller value (10) — it must evict 50; + // 60 > 40 so it's dropped. pk=2 appears mid-stream. + state.insert_batch(&pk_val_batch(&schema, vec![1, 2, 1], vec![10, 99, 60])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 10 |", + "| 1 | 40 |", + "| 2 | 99 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Empty input must produce an empty output stream, not panic. + #[tokio::test] + async fn test_partitioned_topk_rank_empty_input() -> Result<()> { + let (_schema, state) = build_partitioned_topk_rank(3)?; + let results: Vec<_> = state.emit()?.try_collect().await?; + assert!(results.is_empty(), "empty input → empty output"); + Ok(()) + } + + /// `fetch = 1` is a common case (rk = 1 filter) and exercises the + /// boundary-defined-immediately path: after the first admission per + /// partition, `heap.max()` is `Some`, so every subsequent row goes + /// through full Equal/Greater/Less classification. + #[tokio::test] + async fn test_partitioned_topk_rank_fetch_one() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(1)?; + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 1, 2, 2, 3], + vec![3, 1, 9, 4, 7], + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 1 |", + "| 2 | 4 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ORDER BY val DESC` exercises the shared encoder's sort-direction + /// handling: the row converter flips the sort sign for `val` so + /// larger values compare smaller in row-encoded form. Each + /// partition keeps its top-K *largest* values. + #[tokio::test] + async fn test_partitioned_topk_rank_desc_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank_with_opts( + 2, + SortOptions { + descending: true, + nulls_first: false, + }, + false, + )?; + + // pk=1 vals: 10, 5, 8, 12 → top-2 DESC = [12, 10] + // pk=2 vals: 20, 15, 25 → top-2 DESC = [25, 20] + let batch = pk_val_batch( + &schema, + vec![1, 2, 1, 2, 1, 1, 2], + vec![10, 20, 5, 15, 8, 12, 25], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 12 |", + "| 1 | 10 |", + "| 2 | 25 |", + "| 2 | 20 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// NULL sort values exercise the shared encoder's null-ordering + /// handling. With `ASC NULLS LAST`, NULLs sort *after* every + /// non-NULL value, so a partition whose only non-NULL value beats + /// a NULL must evict the NULL when `K = 1`. A partition that holds + /// only NULLs must still emit them. + #[tokio::test] + async fn test_partitioned_topk_rank_nulls_last_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank_with_opts( + 1, + SortOptions { + descending: false, + nulls_first: false, + }, + true, + )?; + + // pk=1 vals: NULL, 7, NULL → top-1 ASC NULLS LAST = [7] + // pk=2 vals: NULL → top-1 = [NULL] + // pk=3 vals: NULL, 4, 2 → top-1 = [2] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 1, 3, 3, 3], + vec![None, None, Some(7), None, None, Some(4), Some(2)], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 7 |", + "| 2 | |", + "| 3 | 2 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ASC NULLS FIRST` (the `SortOptions::default()`) sorts NULLs + /// *before* every non-NULL value, so under `fetch = K` a partition's + /// NULLs are kept preferentially over larger non-NULL values. + #[tokio::test] + async fn test_partitioned_topk_rank_nulls_first_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank_with_opts( + 2, + SortOptions { + descending: false, + nulls_first: true, + }, + true, + )?; + + // pk=1 vals: NULL, 5, NULL, 8 → top-2 ASC NULLS FIRST = [NULL, NULL] + // pk=2 vals: 7, NULL → top-2 = [NULL, 7] + // pk=3 vals: 3, 1 → top-2 = [1, 3] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 3, 1, 2, 1, 3], + vec![ + None, + Some(7), + Some(5), + Some(3), + None, + None, + Some(8), + Some(1), + ], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | |", + "| 1 | |", + "| 2 | |", + "| 2 | 7 |", + "| 3 | 1 |", + "| 3 | 3 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// RANK-specific: heap fills with K rows tied at the same OB value, + /// then more rows at that same value arrive. They take the Equal arm + /// (heap is full, `heap.max() == row`) and accumulate as ties, while + /// strictly-greater rows are dropped. All retained rows have rank 1. + #[tokio::test] + async fn test_partitioned_topk_rank_boundary_ties_retained() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 5, 5, 10, 5 + // - first two 5s fill the heap (max=None until heap reaches K=2) + // - third row 10 > 5 → drop (Greater) + // - fourth row 5 == 5 → push to ties (Equal) + // Sorted RANKs: 5→1, 5→1, 5→1, 10→4. WHERE rk ≤ 2 keeps the three 5s. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 5, 10, 5])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 5 |", + "| 1 | 5 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// RANK-specific: heap fills with K rows tied at value V, equal_indices + /// accumulate at V, then a strictly-better row arrives whose admission + /// shifts the boundary strictly below V. The boundary-changed branch + /// must clear both `state.ties` and the in-flight `equal_indices` — + /// otherwise the now-rank-> K rows at value V would leak into output. + #[tokio::test] + async fn test_partitioned_topk_rank_boundary_shifts_clears_ties() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 10, 10, 5, 3 + // - first two 10s fill heap (max=10) + // - third 10 → Equal → equal_indices=[2] + // - 5 < 10 → admit, evict 10 → heap={5,10}, max=10 (unchanged). + // Push evicted to ties: ties=[10@curr_batch[ev_idx]]. + // - 3 < 10 → admit, evict 10 → heap={3,5}, max=5 (CHANGED). + // Clear ties AND equal_indices. + // Sorted RANKs: 3→1, 5→2, 10→3, 10→3, 10→3. WHERE rk ≤ 2 → [3, 5]. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1, 1], vec![10, 10, 10, 5, 3])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 3 |", + "| 1 | 5 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// RANK-specific: heap has multiple rows at boundary value V, then a + /// strictly-better row arrives. The heap evicts one V (popping + /// `prev_min`), but `heap.max()` is still V — boundary unchanged. + /// The evicted V row must be pushed as a `TieEntry`; without that + /// branch a `rk <= K` query would silently lose a tied row. + #[tokio::test] + async fn test_partitioned_topk_rank_eviction_at_unchanged_boundary() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 10, 5 + // - first two 10s fill the heap (max=10) + // - 5 < 10 → admit, evict 10. New heap={5,10}, max=10 (unchanged). + // Push the evicted 10 to ties. + // Sorted RANKs: 5→1, 10→2, 10→2. WHERE rk ≤ 2 → all 3 rows. + let batch = pk_val_batch(&schema, vec![1, 1, 1], vec![10, 10, 5])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 10 |", + "| 1 | 10 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/tree_node.rs b/native/vendor/datafusion-physical-plan/src/tree_node.rs new file mode 100644 index 0000000..dcdceff --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/tree_node.rs
@@ -0,0 +1,118 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! This module provides common traits for visiting or rewriting tree nodes easily. + +use std::fmt::{self, Display, Formatter}; +use std::sync::Arc; + +use crate::execution_plan::replace_children_if_necessary; +use crate::{ExecutionPlan, displayable}; + +use datafusion_common::Result; +use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode}; + +impl DynTreeNode for dyn ExecutionPlan { + fn arc_children(&self) -> Vec<&Arc<Self>> { + self.children() + } + + fn with_new_arc_children( + &self, + arc_self: Arc<Self>, + new_children: Vec<Arc<Self>>, + ) -> Result<Arc<Self>> { + replace_children_if_necessary(arc_self, new_children) + } +} + +/// A node context object beneficial for writing optimizer rules. +/// This context encapsulating an [`ExecutionPlan`] node with a payload. +/// +/// Since each wrapped node has it's children within both the `PlanContext.plan.children()`, +/// as well as separately within the `PlanContext.children` (which are child nodes wrapped in the context), +/// it's important to keep these child plans in sync when performing mutations. +/// +/// Since there are two ways to access child plans directly -— it's recommended +/// to perform mutable operations via [`Self::update_plan_from_children`]. +/// After mutating the `PlanContext.children`, or after creating the `PlanContext`, +/// call `update_plan_from_children` to sync. +#[derive(Debug)] +pub struct PlanContext<T: Sized> { + /// The execution plan associated with this context. + pub plan: Arc<dyn ExecutionPlan>, + /// Custom data payload of the node. + pub data: T, + /// Child contexts of this node. + pub children: Vec<Self>, +} + +impl<T> PlanContext<T> { + pub fn new(plan: Arc<dyn ExecutionPlan>, data: T, children: Vec<Self>) -> Self { + Self { + plan, + data, + children, + } + } + + /// Update the `PlanContext.plan.children()` from the `PlanContext.children`, + /// if the `PlanContext.children` have been changed. + pub fn update_plan_from_children(mut self) -> Result<Self> { + let children_plans = self.children.iter().map(|c| Arc::clone(&c.plan)).collect(); + self.plan = replace_children_if_necessary(self.plan, children_plans)?; + + Ok(self) + } +} + +impl<T: Default> PlanContext<T> { + pub fn new_default(plan: Arc<dyn ExecutionPlan>) -> Self { + let children = plan + .children() + .into_iter() + .cloned() + .map(Self::new_default) + .collect(); + Self::new(plan, Default::default(), children) + } +} + +impl<T: Display> Display for PlanContext<T> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let node_string = displayable(self.plan.as_ref()).one_line(); + write!(f, "Node plan: {node_string}")?; + write!(f, "Node data: {}", self.data)?; + write!(f, "") + } +} + +impl<T> ConcreteTreeNode for PlanContext<T> { + fn children(&self) -> &[Self] { + &self.children + } + + fn take_children(mut self) -> (Self, Vec<Self>) { + let children = std::mem::take(&mut self.children); + (self, children) + } + + fn with_new_children(mut self, children: Vec<Self>) -> Result<Self> { + self.children = children; + self.update_plan_from_children() + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/union.rs b/native/vendor/datafusion-physical-plan/src/union.rs new file mode 100644 index 0000000..c1cc5da --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/union.rs
@@ -0,0 +1,1786 @@ +// 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. + +// Some of these functions reference the Postgres documentation +// or implementation to ensure compatibility and are subject to +// the Postgres license. + +//! The Union operator combines multiple inputs with the same schema + +use std::borrow::Borrow; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, + metrics::{ExecutionPlanMetricsSet, MetricsSet}, +}; +use crate::execution_plan::{ + CardinalityEffect, InvariantLevel, boundedness_from_children, + check_default_invariants, emission_type_from_children, +}; +use crate::filter::FilterExec; +use crate::filter_pushdown::{ + ChildPushdownResult, FilterDescription, FilterPushdownPhase, + FilterPushdownPropagation, PushedDown, +}; +use crate::metrics::BaselineMetrics; +use crate::projection::{ProjectionExec, ProjectionExpr, make_with_child}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::ObservedStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; + +use arrow::datatypes::{Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion_common::config::ConfigOptions; +use datafusion_common::stats::NdvFallback; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + Result, assert_or_internal_err, exec_err, internal_datafusion_err, plan_err, +}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::{CastExpr, Column}; +use datafusion_physical_expr::{ + EquivalenceProperties, PhysicalExpr, calculate_union, conjunction, +}; + +use futures::Stream; +use itertools::Itertools; +use log::{debug, trace, warn}; +use tokio::macros::support::thread_rng_n; + +/// Coerces `input`'s output schema to exactly `schema` via a `ProjectionExec` +/// that re-stamps each column with the union's merged field (same +/// `DataType`, but the union's merged nullability/name/metadata), or returns +/// `input` unchanged if its schema already matches. [`UnionExec::try_new`] +/// and [`InterleaveExec::try_new`] call this on every child, so the coercion +/// is visible in the plan tree (e.g. in `EXPLAIN`) instead of happening +/// invisibly inside the union operator's own `execute()`. +/// +/// A column whose `DataType` doesn't already match the union's is a genuine +/// data type mismatch (as opposed to a nullability/name/metadata-only one), +/// and is rejected eagerly here rather than silently cast or deferred to a +/// runtime failure -- this only ever changes a column's declared schema, +/// never its values. +/// +/// Casting a column to its own `DataType` (only the `Field`'s nullability, +/// name, or metadata changes) is a zero-copy relabeling: the cast kernel's +/// same-type fast path (`cast_array_by_name`) just clones the `Arc<dyn +/// Array>`, so this carries no runtime overhead over the schema it replaces. +/// +/// See <https://github.com/apache/datafusion/issues/15394>. +fn coerce_schema( + input: Arc<dyn ExecutionPlan>, + schema: &SchemaRef, +) -> Result<Arc<dyn ExecutionPlan>> { + let input_schema = input.schema(); + if &input_schema == schema { + return Ok(input); + } + + let exprs = input_schema + .fields() + .iter() + .zip(schema.fields()) + .enumerate() + .map(|(i, (input_field, target_field))| { + if input_field.data_type() != target_field.data_type() { + return plan_err!( + "UnionExec/InterleaveExec requires all inputs to have the same \ + data type per column; column {i} has type {} in one input, but \ + the union schema expects {}", + input_field.data_type(), + target_field.data_type() + ); + } + let column: Arc<dyn PhysicalExpr> = + Arc::new(Column::new(input_field.name(), i)); + let expr = if input_field == target_field { + column + } else { + Arc::new(CastExpr::new_with_target_field( + column, + Arc::clone(target_field), + None, + )) as Arc<dyn PhysicalExpr> + }; + Ok(ProjectionExpr { + expr, + alias: target_field.name().clone(), + }) + }) + .collect::<Result<Vec<_>>>()?; + + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) +} + +/// `UnionExec`: `UNION ALL` execution plan. +/// +/// `UnionExec` combines multiple inputs with the same schema by +/// concatenating the partitions. It does not mix or copy data within +/// or across partitions. Thus if the input partitions are sorted, the +/// output partitions of the union are also sorted. +/// +/// For example, given a `UnionExec` of two inputs, with `N` +/// partitions, and `M` partitions, there will be `N+M` output +/// partitions. The first `N` output partitions are from Input 1 +/// partitions, and then next `M` output partitions are from Input 2. +/// +/// ```text +/// ▲ ▲ ▲ ▲ +/// │ │ │ │ +/// Output │ ... │ │ │ +/// Partitions │0 │N-1 │ N │N+M-1 +/// (passes through ┌────┴───────┴───────────┴─────────┴───┐ +/// the N+M input │ UnionExec │ +/// partitions) │ │ +/// └──────────────────────────────────────┘ +/// ▲ +/// │ +/// │ +/// Input ┌────────┬─────┴────┬──────────┐ +/// Partitions │ ... │ │ ... │ +/// 0 │ │ N-1 │ 0 │ M-1 +/// ┌────┴────────┴───┐ ┌───┴──────────┴───┐ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │ │ │ │ +/// │Input 1 │ │Input 2 │ +/// └─────────────────┘ └──────────────────┘ +/// ``` +#[derive(Debug, Clone)] +pub struct UnionExec { + /// Input execution plan + inputs: Vec<Arc<dyn ExecutionPlan>>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl UnionExec { + /// Try to create a new UnionExec. + /// + /// # Errors + /// Returns an error if: + /// - `inputs` is empty + /// + /// # Optimization + /// If there is only one input, returns that input directly rather than wrapping it in a UnionExec + pub fn try_new( + inputs: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + match inputs.len() { + 0 => exec_err!("UnionExec requires at least one input"), + 1 => Ok(inputs.into_iter().next().unwrap()), + _ => { + let schema = union_schema(&inputs)?; + // The schema of the inputs and the union schema is consistent when: + // - They have the same number of fields, and + // - Their fields have same types at the same indices. + let inputs = inputs + .into_iter() + .map(|input| coerce_schema(input, &schema)) + .collect::<Result<Vec<_>>>()?; + let cache = Self::compute_properties(&inputs, schema)?; + Ok(Arc::new(UnionExec { + inputs, + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + })) + } + } + } + + /// Get inputs of the execution plan + pub fn inputs(&self) -> &Vec<Arc<dyn ExecutionPlan>> { + &self.inputs + } + + /// Maps a global output partition index to the `(input index, local + /// partition index)` of the input that owns it, or `None` if out of range. + fn owning_input(&self, partition: usize) -> Option<(usize, usize)> { + let mut remaining = partition; + for (i, input) in self.inputs.iter().enumerate() { + let count = input.output_partitioning().partition_count(); + if remaining < count { + return Some((i, remaining)); + } + remaining -= count; + } + None + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + inputs: &[Arc<dyn ExecutionPlan>], + schema: SchemaRef, + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let children_eqps = inputs + .iter() + .map(|child| child.equivalence_properties().clone()) + .collect::<Vec<_>>(); + let eq_properties = calculate_union(children_eqps, schema)?; + + // Calculate output partitioning; i.e. sum output partitions of the inputs. + let num_partitions = inputs + .iter() + .map(|plan| plan.output_partitioning().partition_count()) + .sum(); + let output_partitioning = Partitioning::UnknownPartitioning(num_partitions); + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + emission_type_from_children(inputs), + boundedness_from_children(inputs), + )) + } +} + +impl DisplayAs for UnionExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "UnionExec") + } + DisplayFormatType::TreeRender => Ok(()), + } + } +} + +impl ExecutionPlan for UnionExec { + fn name(&self) -> &'static str { + "UnionExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn check_invariants(&self, check: InvariantLevel) -> Result<()> { + check_default_invariants(self, check)?; + + (self.inputs().len() >= 2).then_some(()).ok_or_else(|| { + internal_datafusion_err!("UnionExec should have at least 2 children") + }) + } + + fn maintains_input_order(&self) -> Vec<bool> { + // If the Union has an output ordering, it maintains at least one + // child's ordering (i.e. the meet). + // For instance, assume that the first child is SortExpr('a','b','c'), + // the second child is SortExpr('a','b') and the third child is + // SortExpr('a','b'). The output ordering would be SortExpr('a','b'), + // which is the "meet" of all input orderings. In this example, this + // function will return vec![false, true, true], indicating that we + // preserve the orderings for the 2nd and the 3rd children. + if let Some(output_ordering) = self.properties().output_ordering() { + self.inputs() + .iter() + .map(|child| { + if let Some(child_ordering) = child.output_ordering() { + output_ordering.len() == child_ordering.len() + } else { + false + } + }) + .collect() + } else { + vec![false; self.inputs().len()] + } + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false; self.children().len()] + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + self.inputs.iter().collect() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => UnionExec::try_new(children), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + mut partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start UnionExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + // record the tiny amount of work done in this function so + // elapsed_compute is reported as non zero + let elapsed_compute = baseline_metrics.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); // record on drop + + // find partition to execute + for input in self.inputs.iter() { + // Calculate whether partition belongs to the current partition + if partition < input.output_partitioning().partition_count() { + let stream = input.execute(partition, context)?; + debug!("Found a Union partition to execute"); + return Ok(Box::pin(ObservedStream::new( + stream, + baseline_metrics, + None, + ))); + } else { + partition -= input.output_partitioning().partition_count(); + } + } + + warn!("Error in Union: Partition {partition} not found"); + + exec_err!("Partition {partition} not found in Union") + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + if let Some(partition_idx) = partition { + // For a specific partition, compute stats only for the input that + // owns it; the other inputs are not needed and are skipped. + let targeted = self.owning_input(partition_idx); + self.inputs + .iter() + .enumerate() + .map(|(i, _)| match targeted { + Some((target_i, target_partition)) if i == target_i => { + ChildStats::At(Some(target_partition)) + } + _ => ChildStats::Skip, + }) + .collect() + } else { + vec![ChildStats::At(None); self.inputs.len()] + } + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + if let Some(partition_idx) = args.partition() { + // For a specific partition, find which input it belongs to + if let Some((target_i, _)) = self.owning_input(partition_idx) { + // This partition belongs to this input - return its stats + return Ok(Arc::clone(&input_stats[target_i])); + } + // If we get here, the partition index is out of bounds + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } else { + let stats_refs = input_stats.iter().map(|s| s.as_ref()).collect::<Vec<_>>(); + + Ok(Arc::new(Statistics::try_merge_iter_with_ndv_fallback( + stats_refs, + self.schema().as_ref(), + NdvFallback::Sum, + )?)) + } + } + + fn cardinality_effect(&self) -> CardinalityEffect { + // Union combines rows from multiple inputs, so output rows are not tied + // to any single input and can only be constrained as greater-or-equal. + CardinalityEffect::GreaterEqual + } + + fn supports_limit_pushdown(&self) -> bool { + true + } + + /// Tries to push `projection` down through `union`. If possible, performs the + /// pushdown and returns a new [`UnionExec`] as the top plan which has projections + /// as its children. Otherwise, returns `None`. + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // If the projection doesn't narrow the schema, we shouldn't try to push it down. + if projection.expr().len() >= projection.input().schema().fields().len() { + return Ok(None); + } + + let new_children = self + .children() + .into_iter() + .map(|child| make_with_child(projection, child)) + .collect::<Result<Vec<_>>>()?; + + Ok(Some(UnionExec::try_new(new_children.clone())?)) + } + + fn gather_filters_for_pushdown( + &self, + _phase: FilterPushdownPhase, + parent_filters: Vec<Arc<dyn PhysicalExpr>>, + _config: &ConfigOptions, + ) -> Result<FilterDescription> { + FilterDescription::from_children(parent_filters, &self.children()) + } + + fn handle_child_pushdown_result( + &self, + phase: FilterPushdownPhase, + child_pushdown_result: ChildPushdownResult, + _config: &ConfigOptions, + ) -> Result<FilterPushdownPropagation<Arc<dyn ExecutionPlan>>> { + // Pre phase: handle heterogeneous pushdown by wrapping individual + // children with FilterExec and reporting all filters as handled. + // Post phase: use default behavior to let the filter creator decide how to handle + // filters that weren't fully pushed down. + if phase != FilterPushdownPhase::Pre { + return Ok(FilterPushdownPropagation::if_all(child_pushdown_result)); + } + + // UnionExec needs specialized filter pushdown handling when children have + // heterogeneous pushdown support. Without this, when some children support + // pushdown and others don't, the default behavior would leave FilterExec + // above UnionExec, re-applying filters to outputs of all children—including + // those that already applied the filters via pushdown. This specialized + // implementation adds FilterExec only to children that don't support + // pushdown, avoiding redundant filtering and improving performance. + // + // Example: Given Child1 (no pushdown support) and Child2 (has pushdown support) + // Default behavior: This implementation: + // FilterExec UnionExec + // UnionExec FilterExec + // Child1 Child1 + // Child2(filter) Child2(filter) + + // Collect unsupported filters for each child + let mut unsupported_filters_per_child = vec![Vec::new(); self.inputs.len()]; + for parent_filter_result in child_pushdown_result.parent_filters.iter() { + for (child_idx, &child_result) in + parent_filter_result.child_results.iter().enumerate() + { + if matches!(child_result, PushedDown::No) { + unsupported_filters_per_child[child_idx] + .push(Arc::clone(&parent_filter_result.filter)); + } + } + } + + // Wrap children that have unsupported filters with FilterExec + let mut new_children = self.inputs.clone(); + for (child_idx, unsupported_filters) in + unsupported_filters_per_child.iter().enumerate() + { + if !unsupported_filters.is_empty() { + let combined_filter = conjunction(unsupported_filters.clone()); + new_children[child_idx] = Arc::new(FilterExec::try_new( + combined_filter, + Arc::clone(&self.inputs[child_idx]), + )?); + } + } + + // Check if any children were modified + let children_modified = new_children + .iter() + .zip(self.inputs.iter()) + .any(|(new, old)| !Arc::ptr_eq(new, old)); + + let all_filters_pushed = + vec![PushedDown::Yes; child_pushdown_result.parent_filters.len()]; + let propagation = if children_modified { + let updated_node = UnionExec::try_new(new_children)?; + FilterPushdownPropagation::with_parent_pushdown_result(all_filters_pushed) + .with_updated_node(updated_node) + } else { + FilterPushdownPropagation::with_parent_pushdown_result(all_filters_pushed) + }; + + // Report all parent filters as supported since we've ensured they're applied + // on all children (either pushed down or via FilterExec) + Ok(propagation) + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Union( + protobuf::UnionExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnionExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let union = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Union, + "UnionExec", + ); + let inputs = union + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::<Result<Vec<_>>>()?; + UnionExec::try_new(inputs) + } +} + +/// Combines multiple input streams by interleaving them. +/// +/// All inputs must share an identical [`Partitioning::Hash`] or [`Partitioning::Range`] so that +/// partition `k` covers the same data across every input. Each output partition is the +/// interleaving of the same-indexed partition from all inputs: +/// `output[k] = input[0][k] + input[1][k] + ... + input[n-1][k]` +/// +/// # Data Flow +/// ```text +/// +---------+ +/// | |---+ +/// | Input 1 | | +/// | |-------------+ +/// +---------+ | | +/// | | +---------+ +/// +------------------>| | +/// +---------------->| Combine |--> +/// | +-------------->| | +/// | | | +---------+ +/// +---------+ | | | +/// | |-----+ | | +/// | Input 2 | | | +/// | |---------------+ +/// +---------+ | | | +/// | | | +---------+ +/// | +-------->| | +/// | +------>| Combine |--> +/// | +---->| | +/// | | +---------+ +/// +---------+ | | +/// | |-------+ | +/// | Input 3 | | +/// | |-----------------+ +/// +---------+ +/// ``` +#[derive(Debug, Clone)] +pub struct InterleaveExec { + /// Input execution plan + inputs: Vec<Arc<dyn ExecutionPlan>>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl InterleaveExec { + /// Create a new InterleaveExec + pub fn try_new(inputs: Vec<Arc<dyn ExecutionPlan>>) -> Result<Self> { + assert_or_internal_err!( + can_interleave(inputs.iter()), + "Not all InterleaveExec children have a consistent hash or range partitioning" + ); + let schema = union_schema(&inputs)?; + let inputs = inputs + .into_iter() + .map(|input| coerce_schema(input, &schema)) + .collect::<Result<Vec<_>>>()?; + let cache = Self::compute_properties(&inputs, schema)?; + Ok(InterleaveExec { + inputs, + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + }) + } + + /// Get inputs of the execution plan + pub fn inputs(&self) -> &Vec<Arc<dyn ExecutionPlan>> { + &self.inputs + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + inputs: &[Arc<dyn ExecutionPlan>], + schema: SchemaRef, + ) -> Result<PlanProperties> { + let eq_properties = EquivalenceProperties::new(schema); + // Get output partitioning: + let output_partitioning = inputs[0].output_partitioning().clone(); + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + emission_type_from_children(inputs), + boundedness_from_children(inputs), + )) + } +} + +impl DisplayAs for InterleaveExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "InterleaveExec") + } + DisplayFormatType::TreeRender => Ok(()), + } + } +} + +impl ExecutionPlan for InterleaveExec { + fn name(&self) -> &'static str { + "InterleaveExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + self.inputs.iter().collect() + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![false; self.inputs().len()] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + // New children are no longer interleavable, which might be a bug of optimization rewrite. + assert_or_internal_err!( + can_interleave(children.iter()), + "Can not create InterleaveExec: new children can not be interleaved" + ); + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + trace!( + "Start InterleaveExec::execute for partition {} of context session_id {} and task_id {:?}", + partition, + context.session_id(), + context.task_id() + ); + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + // record the tiny amount of work done in this function so + // elapsed_compute is reported as non zero + let elapsed_compute = baseline_metrics.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); // record on drop + + let mut input_stream_vec = vec![]; + for input in self.inputs.iter() { + if partition < input.output_partitioning().partition_count() { + let stream = input.execute(partition, Arc::clone(&context))?; + input_stream_vec.push(stream); + } else { + // Do not find a partition to execute + break; + } + } + if input_stream_vec.len() == self.inputs.len() { + let stream = Box::pin(CombinedRecordBatchStream::new( + self.schema(), + input_stream_vec, + )); + return Ok(Box::pin(ObservedStream::new( + stream, + baseline_metrics, + None, + ))); + } + + warn!("Error in InterleaveExec: Partition {partition} not found"); + + exec_err!("Partition {partition} not found in InterleaveExec") + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition); self.inputs.len()] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let stats = input_stats + .iter() + .map(|s| s.as_ref().clone()) + .collect::<Vec<_>>(); + + Ok(Arc::new(Statistics::try_merge_iter_with_ndv_fallback( + stats.iter(), + self.schema().as_ref(), + NdvFallback::Sum, + )?)) + } + + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false; self.children().len()] + } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Interleave( + protobuf::InterleaveExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl InterleaveExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + let interleave = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Interleave, + "InterleaveExec", + ); + let inputs = interleave + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::<Result<Vec<_>>>()?; + Ok(Arc::new(InterleaveExec::try_new(inputs)?)) + } +} + +/// Returns true if all inputs have the same [`Partitioning::Hash`] or [`Partitioning::Range`] +/// spec, making them safe to interleave. Two inputs are interleave-compatible when partition +/// `k` covers the identical key range or hash bucket across every input. +/// +/// Note: compatibility is checked sequentially against the first input, so +/// `InputDistributionRequirements::co_partitioned` is not needed here. +/// +/// It might be too strict here in the case that the input partition specs are compatible but not exactly the same. +/// For example one input partition has the partition spec Hash('a','b','c') and +/// other has the partition spec Hash('a'), It is safe to derive the out partition with the spec Hash('a','b','c'). +pub fn can_interleave<T: Borrow<Arc<dyn ExecutionPlan>>>( + mut inputs: impl Iterator<Item = T>, +) -> bool { + let Some(first) = inputs.next() else { + return false; + }; + + let reference = first.borrow().output_partitioning(); + matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_)) + && inputs + .map(|plan| plan.borrow().output_partitioning().clone()) + .all(|partition| partition == *reference) +} + +fn union_schema(inputs: &[Arc<dyn ExecutionPlan>]) -> Result<SchemaRef> { + if inputs.is_empty() { + return exec_err!("Cannot create union schema from empty inputs"); + } + + let first_schema = inputs[0].schema(); + let first_field_count = first_schema.fields().len(); + + // validate that all inputs have the same number of fields + for (idx, input) in inputs.iter().enumerate().skip(1) { + let field_count = input.schema().fields().len(); + if field_count != first_field_count { + return exec_err!( + "UnionExec/InterleaveExec requires all inputs to have the same number of fields. \ + Input 0 has {first_field_count} fields, but input {idx} has {field_count} fields" + ); + } + } + + let fields = (0..first_field_count) + .map(|i| { + // We take the name from the left side of the union to match how names are coerced during logical planning, + // which also uses the left side names. + let base_field = first_schema.field(i).clone(); + + // Coerce metadata and nullability across all inputs + + inputs + .iter() + .enumerate() + .map(|(input_idx, input)| { + let field = input.schema().field(i).clone(); + let mut metadata = field.metadata().clone(); + + let other_metadatas = inputs + .iter() + .enumerate() + .filter(|(other_idx, _)| *other_idx != input_idx) + .flat_map(|(_, other_input)| { + other_input.schema().field(i).metadata().clone().into_iter() + }); + + metadata.extend(other_metadatas); + field.with_metadata(metadata) + }) + .find_or_first(Field::is_nullable) + // We can unwrap this because if inputs was empty, this would've already panic'ed when we + // indexed into inputs[0]. + .unwrap() + .with_name(base_field.name()) + }) + .collect::<Vec<_>>(); + + let all_metadata_merged = inputs + .iter() + .flat_map(|i| i.schema().metadata().clone().into_iter()) + .collect(); + + Ok(Arc::new(Schema::new_with_metadata( + fields, + all_metadata_merged, + ))) +} + +/// CombinedRecordBatchStream can be used to combine a Vec of SendableRecordBatchStreams into one +struct CombinedRecordBatchStream { + /// Schema wrapped by Arc + schema: SchemaRef, + /// Stream entries + entries: Vec<SendableRecordBatchStream>, +} + +impl CombinedRecordBatchStream { + /// Create an CombinedRecordBatchStream + pub fn new(schema: SchemaRef, entries: Vec<SendableRecordBatchStream>) -> Self { + Self { schema, entries } + } +} + +impl RecordBatchStream for CombinedRecordBatchStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for CombinedRecordBatchStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + use Poll::*; + + let start = thread_rng_n(self.entries.len() as u32) as usize; + let mut idx = start; + + for _ in 0..self.entries.len() { + let stream = self.entries.get_mut(idx).unwrap(); + + match Pin::new(stream).poll_next(cx) { + Ready(Some(val)) => return Ready(Some(val)), + Ready(None) => { + // Remove the entry + self.entries.swap_remove(idx); + + // Check if this was the last entry, if so the cursor needs + // to wrap + if idx == self.entries.len() { + idx = 0; + } else if idx < start && start <= self.entries.len() { + // The stream being swapped into the current index has + // already been polled, so skip it. + idx = idx.wrapping_add(1) % self.entries.len(); + } + } + Pending => { + idx = idx.wrapping_add(1) % self.entries.len(); + } + } + } + + // If the map is empty, then the stream is complete. + if self.entries.is_empty() { + Ready(None) + } else { + Pending + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect; + use crate::repartition::RepartitionExec; + use crate::statistics::{StatisticsArgs, StatisticsContext}; + use crate::test::exec::StatisticsExec; + use crate::test::{self, TestMemoryExec}; + + use arrow::compute::SortOptions; + use arrow::datatypes::DataType; + use datafusion_common::SplitPoint; + use datafusion_common::stats::Precision; + use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_physical_expr::RangePartitioning; + use datafusion_physical_expr::equivalence::convert_to_orderings; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; + + // Generate a schema which consists of 7 columns (a, b, c, d, e, f, g) + fn create_test_schema() -> Result<SchemaRef> { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, true); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, true); + let e = Field::new("e", DataType::Int32, true); + let f = Field::new("f", DataType::Int32, true); + let g = Field::new("g", DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e, f, g])); + + Ok(schema) + } + + fn create_test_schema2() -> Result<SchemaRef> { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, true); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, true); + let e = Field::new("e", DataType::Int32, true); + let f = Field::new("f", DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e, f])); + + Ok(schema) + } + + #[tokio::test] + async fn test_union_partitions() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + + // Create inputs with different partitioning + let csv = test::scan_partitioned(4); + let csv2 = test::scan_partitioned(5); + + let union_exec: Arc<dyn ExecutionPlan> = UnionExec::try_new(vec![csv, csv2])?; + + // Should have 9 partitions and 9 output batches + assert_eq!( + union_exec + .properties() + .output_partitioning() + .partition_count(), + 9 + ); + + let result: Vec<RecordBatch> = collect(union_exec, task_ctx).await?; + assert_eq!(result.len(), 9); + + Ok(()) + } + + #[tokio::test] + async fn test_interleave_conforms_batch_schema() -> Result<()> { + // Two inputs agree on the column's type but disagree on nullability; + // InterleaveExec's declared schema ORs nullability across inputs, so + // every yielded batch must be re-stamped with that schema. See + // <https://github.com/apache/datafusion/issues/15394>. + let task_ctx = Arc::new(TaskContext::default()); + + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch_not_null = RecordBatch::try_new( + Arc::clone(&schema_not_null), + vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2]))], + )?; + + let schema_nullable = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch_nullable = RecordBatch::try_new( + Arc::clone(&schema_nullable), + vec![Arc::new(arrow::array::Int32Array::from(vec![3, 4]))], + )?; + + let hash_expr = vec![col("a", schema_not_null.as_ref())?]; + let left: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new( + TestMemoryExec::try_new_exec(&[vec![batch_not_null]], schema_not_null, None)?, + Partitioning::Hash(hash_expr.clone(), 1), + )?); + let right: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new( + TestMemoryExec::try_new_exec(&[vec![batch_nullable]], schema_nullable, None)?, + Partitioning::Hash(hash_expr, 1), + )?); + + let interleave: Arc<dyn ExecutionPlan> = + Arc::new(InterleaveExec::try_new(vec![left, right])?); + let interleave_schema = interleave.schema(); + assert!(interleave_schema.field(0).is_nullable()); + + let batches = collect(interleave, task_ctx).await?; + assert!(!batches.is_empty()); + for batch in &batches { + assert_eq!(batch.schema(), interleave_schema); + } + + Ok(()) + } + + fn stats_merge_inputs() -> (SchemaRef, Statistics, Statistics, Statistics) { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])); + + let left = Statistics::default() + .with_num_rows(Precision::Exact(5)) + .with_total_byte_size(Precision::Exact(23)) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Exact(5)) + .with_min_value(Precision::Exact(ScalarValue::UInt32(Some(1)))) + .with_max_value(Precision::Exact(ScalarValue::UInt32(Some(21)))) + .with_sum_value(Precision::Exact(ScalarValue::UInt32(Some(42)))) + .with_null_count(Precision::Exact(0)) + .with_byte_size(Precision::Exact(40)), + ); + + let right = Statistics::default() + .with_num_rows(Precision::Exact(7)) + .with_total_byte_size(Precision::Exact(29)) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Exact(3)) + .with_min_value(Precision::Exact(ScalarValue::UInt32(Some(22)))) + .with_max_value(Precision::Exact(ScalarValue::UInt32(Some(34)))) + .with_sum_value(Precision::Exact(ScalarValue::UInt32(Some(8)))) + .with_null_count(Precision::Exact(1)) + .with_byte_size(Precision::Exact(60)), + ); + + let expected = Statistics::default() + .with_num_rows(Precision::Exact(12)) + .with_total_byte_size(Precision::Exact(52)) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Inexact(8)) + .with_min_value(Precision::Exact(ScalarValue::UInt32(Some(1)))) + .with_max_value(Precision::Exact(ScalarValue::UInt32(Some(34)))) + .with_sum_value(Precision::Exact(ScalarValue::UInt64(Some(50)))) + .with_null_count(Precision::Exact(1)) + .with_byte_size(Precision::Exact(100)), + ); + + (schema, left, right, expected) + } + + fn stats_merge_multicolumn_inputs() -> (SchemaRef, Statistics, Statistics, Statistics) + { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Float32, true), + ])); + + let left = Statistics::default() + .with_num_rows(Precision::Exact(5)) + .with_total_byte_size(Precision::Exact(23)) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Exact(5)) + .with_min_value(Precision::Exact(ScalarValue::Int64(Some(-4)))) + .with_max_value(Precision::Exact(ScalarValue::Int64(Some(21)))) + .with_sum_value(Precision::Exact(ScalarValue::Int64(Some(42)))) + .with_null_count(Precision::Exact(0)), + ) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Exact(2)) + .with_min_value(Precision::Exact(ScalarValue::from("a"))) + .with_max_value(Precision::Exact(ScalarValue::from("x"))) + .with_null_count(Precision::Exact(3)), + ) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_max_value(Precision::Exact(ScalarValue::Float32(Some(1.1)))) + .with_min_value(Precision::Exact(ScalarValue::Float32(Some(0.1)))) + .with_sum_value(Precision::Exact(ScalarValue::Float32(Some(42.0)))), + ); + + let right = Statistics::default() + .with_num_rows(Precision::Exact(7)) + .with_total_byte_size(Precision::Exact(29)) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Exact(3)) + .with_min_value(Precision::Exact(ScalarValue::Int64(Some(1)))) + .with_max_value(Precision::Exact(ScalarValue::Int64(Some(34)))) + .with_sum_value(Precision::Exact(ScalarValue::Int64(Some(42)))) + .with_null_count(Precision::Exact(1)), + ) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Exact(3)) + .with_min_value(Precision::Exact(ScalarValue::from("b"))) + .with_max_value(Precision::Exact(ScalarValue::from("z"))), + ) + .add_column_statistics(ColumnStatistics::new_unknown()); + + let expected = Statistics::default() + .with_num_rows(Precision::Exact(12)) + .with_total_byte_size(Precision::Exact(52)) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Inexact(6)) + .with_min_value(Precision::Exact(ScalarValue::Int64(Some(-4)))) + .with_max_value(Precision::Exact(ScalarValue::Int64(Some(34)))) + .with_sum_value(Precision::Exact(ScalarValue::Int64(Some(84)))) + .with_null_count(Precision::Exact(1)), + ) + .add_column_statistics( + ColumnStatistics::new_unknown() + .with_distinct_count(Precision::Inexact(5)) + .with_min_value(Precision::Exact(ScalarValue::from("a"))) + .with_max_value(Precision::Exact(ScalarValue::from("z"))), + ) + .add_column_statistics(ColumnStatistics::new_unknown()); + + (schema, left, right, expected) + } + + #[test] + fn test_union_partition_statistics_uses_shared_statistics_merge() -> Result<()> { + let (schema, left, right, expected) = stats_merge_inputs(); + + let left: Arc<dyn ExecutionPlan> = + Arc::new(StatisticsExec::new(left, schema.as_ref().clone())); + let right: Arc<dyn ExecutionPlan> = + Arc::new(StatisticsExec::new(right, schema.as_ref().clone())); + + let union = UnionExec::try_new(vec![left, right])?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[test] + fn test_union_partition_statistics_uses_shared_statistics_merge_multicolumn() + -> Result<()> { + let (schema, left, right, expected) = stats_merge_multicolumn_inputs(); + + let left: Arc<dyn ExecutionPlan> = + Arc::new(StatisticsExec::new(left, schema.as_ref().clone())); + let right: Arc<dyn ExecutionPlan> = + Arc::new(StatisticsExec::new(right, schema.as_ref().clone())); + + let union = UnionExec::try_new(vec![left, right])?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[test] + fn test_union_partition_statistics_with_mismatched_nullability() -> Result<()> { + // Regression test for the `ProjectionExec` wrapper `UnionExec::try_new` + // inserts above the non-nullable leg here (via `coerce_schema`): + // exact column statistics (min/max/null/distinct/sum/byte_size) must + // still make it through the wrapper's same-type `CastExpr`, not get + // poisoned into `Absent` the way a generic (type-changing) cast's + // statistics would be. + let (_, left, right, expected) = stats_merge_inputs(); + + // `total_byte_size` differs from the plain-merge fixture (52): the + // wrapper is a `ProjectionExec`, whose `statistics_from_inputs` + // recomputes `total_byte_size` from the (unchanged) schema's row + // width times row count, rather than trusting the wrapped leg's own + // self-reported total -- still `Exact`, just derived differently. + // left: 5 rows * 4 bytes (UInt32) = 20 (was 23); right is untouched + // (already nullable, so `coerce_schema` doesn't wrap it): 20 + 29 = 49. + let expected = expected.with_total_byte_size(Precision::Exact(49)); + + let non_nullable_schema = + Schema::new(vec![Field::new("a", DataType::UInt32, false)]); + let nullable_schema = Schema::new(vec![Field::new("a", DataType::UInt32, true)]); + + let left: Arc<dyn ExecutionPlan> = + Arc::new(StatisticsExec::new(left, non_nullable_schema)); + let right: Arc<dyn ExecutionPlan> = + Arc::new(StatisticsExec::new(right, nullable_schema)); + + let union = UnionExec::try_new(vec![left, right])?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[tokio::test] + async fn test_coerce_schema_no_op_when_already_matching() -> Result<()> { + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc<dyn ExecutionPlan> = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_not_null), None)?; + + let coerced = coerce_schema(Arc::clone(&input), &schema_not_null)?; + assert!(Arc::ptr_eq(&coerced, &input)); + + Ok(()) + } + + #[tokio::test] + async fn test_coerce_schema_casts_only_nullability() -> Result<()> { + // Mismatched nullability: the input gets wrapped in a `ProjectionExec` + // whose `CastExpr` re-stamps the column with the target's `Field` + // (same `DataType`, so this is a zero-copy relabeling, not a real cast). + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch_not_null = RecordBatch::try_new( + Arc::clone(&schema_not_null), + vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2]))], + )?; + let input: Arc<dyn ExecutionPlan> = TestMemoryExec::try_new_exec( + &[vec![batch_not_null]], + Arc::clone(&schema_not_null), + None, + )?; + + let nullable_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let coerced = coerce_schema(Arc::clone(&input), &nullable_schema)?; + assert_eq!(&coerced.schema(), &nullable_schema); + let plan_str = crate::displayable(coerced.as_ref()) + .indent(true) + .to_string(); + assert!( + plan_str.contains("CAST"), + "expected a CAST in the coerced plan:\n{plan_str}" + ); + + let task_ctx = Arc::new(TaskContext::default()); + let batches = collect(coerced, task_ctx).await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema(), nullable_schema); + + Ok(()) + } + + #[test] + fn test_coerce_schema_rejects_genuine_type_mismatch() -> Result<()> { + let schema_int = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc<dyn ExecutionPlan> = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_int), None)?; + + let schema_utf8 = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)])); + let err = coerce_schema(input, &schema_utf8).unwrap_err(); + assert!(err.to_string().contains("same data type per column")); + + Ok(()) + } + + #[test] + fn test_interleave_partition_statistics_uses_shared_statistics_merge() -> Result<()> { + let (schema, left, right, expected) = stats_merge_inputs(); + let hash_expr = vec![col("a", schema.as_ref())?]; + + let left: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new( + Arc::new(StatisticsExec::new(left, schema.as_ref().clone())), + Partitioning::Hash(hash_expr.clone(), 2), + )?); + let right: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new( + Arc::new(StatisticsExec::new(right, schema.as_ref().clone())), + Partitioning::Hash(hash_expr, 2), + )?); + + let interleave = InterleaveExec::try_new(vec![left, right])?; + let stats = + StatisticsContext::new().compute(&interleave, &StatisticsArgs::new())?; + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[test] + fn test_interleave_partition_statistics_for_partition_uses_shared_statistics_merge() + -> Result<()> { + let (schema, left, right, _) = stats_merge_inputs(); + let hash_expr = vec![col("a", schema.as_ref())?]; + + let left: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new( + Arc::new(StatisticsExec::new(left, schema.as_ref().clone())), + Partitioning::Hash(hash_expr.clone(), 2), + )?); + let right: Arc<dyn ExecutionPlan> = Arc::new(RepartitionExec::try_new( + Arc::new(StatisticsExec::new(right, schema.as_ref().clone())), + Partitioning::Hash(hash_expr, 2), + )?); + + let interleave = InterleaveExec::try_new(vec![left, right])?; + let stats = StatisticsContext::new() + .compute(&interleave, &StatisticsArgs::new().with_partition(Some(0)))?; + + let expected = Statistics::default() + .with_num_rows(Precision::Inexact(5)) + .with_total_byte_size(Precision::Inexact(25)) + .add_column_statistics(ColumnStatistics::new_unknown()); + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[tokio::test] + async fn test_union_equivalence_properties() -> Result<()> { + let schema = create_test_schema()?; + let col_a = &col("a", &schema)?; + let col_b = &col("b", &schema)?; + let col_c = &col("c", &schema)?; + let col_d = &col("d", &schema)?; + let col_e = &col("e", &schema)?; + let col_f = &col("f", &schema)?; + let options = SortOptions::default(); + let test_cases = [ + //-----------TEST CASE 1----------// + ( + // First child orderings + vec![ + // [a ASC, b ASC, f ASC] + vec![(col_a, options), (col_b, options), (col_f, options)], + ], + // Second child orderings + vec![ + // [a ASC, b ASC, c ASC] + vec![(col_a, options), (col_b, options), (col_c, options)], + // [a ASC, b ASC, f ASC] + vec![(col_a, options), (col_b, options), (col_f, options)], + ], + // Union output orderings + vec![ + // [a ASC, b ASC, f ASC] + vec![(col_a, options), (col_b, options), (col_f, options)], + ], + ), + //-----------TEST CASE 2----------// + ( + // First child orderings + vec![ + // [a ASC, b ASC, f ASC] + vec![(col_a, options), (col_b, options), (col_f, options)], + // d ASC + vec![(col_d, options)], + ], + // Second child orderings + vec![ + // [a ASC, b ASC, c ASC] + vec![(col_a, options), (col_b, options), (col_c, options)], + // [e ASC] + vec![(col_e, options)], + ], + // Union output orderings + vec![ + // [a ASC, b ASC] + vec![(col_a, options), (col_b, options)], + ], + ), + ]; + + for ( + test_idx, + (first_child_orderings, second_child_orderings, union_orderings), + ) in test_cases.iter().enumerate() + { + let first_orderings = convert_to_orderings(first_child_orderings); + let second_orderings = convert_to_orderings(second_child_orderings); + let union_expected_orderings = convert_to_orderings(union_orderings); + let child1_exec = TestMemoryExec::try_new(&[], Arc::clone(&schema), None)? + .try_with_sort_information(first_orderings)?; + let child1 = Arc::new(child1_exec); + let child1 = Arc::new(TestMemoryExec::update_cache(&child1)); + let child2_exec = TestMemoryExec::try_new(&[], Arc::clone(&schema), None)? + .try_with_sort_information(second_orderings)?; + let child2 = Arc::new(child2_exec); + let child2 = Arc::new(TestMemoryExec::update_cache(&child2)); + + let mut union_expected_eq = EquivalenceProperties::new(Arc::clone(&schema)); + union_expected_eq.add_orderings(union_expected_orderings); + + let union: Arc<dyn ExecutionPlan> = UnionExec::try_new(vec![child1, child2])?; + let union_eq_properties = union.properties().equivalence_properties(); + let err_msg = format!( + "Error in test id: {:?}, test case: {:?}", + test_idx, test_cases[test_idx] + ); + assert_eq_properties_same(union_eq_properties, &union_expected_eq, err_msg); + } + Ok(()) + } + + fn assert_eq_properties_same( + lhs: &EquivalenceProperties, + rhs: &EquivalenceProperties, + err_msg: String, + ) { + // Check whether orderings are same. + let lhs_orderings = lhs.oeq_class(); + let rhs_orderings = rhs.oeq_class(); + assert_eq!(lhs_orderings.len(), rhs_orderings.len(), "{err_msg}"); + for rhs_ordering in rhs_orderings.iter() { + assert!(lhs_orderings.contains(rhs_ordering), "{}", err_msg); + } + } + + #[test] + fn test_union_empty_inputs() { + // Test that UnionExec::try_new fails with empty inputs + let result = UnionExec::try_new(vec![]); + assert!( + result + .unwrap_err() + .to_string() + .contains("UnionExec requires at least one input") + ); + } + + #[test] + fn test_union_schema_empty_inputs() { + // Test that union_schema fails with empty inputs + let result = union_schema(&[]); + assert!( + result + .unwrap_err() + .to_string() + .contains("Cannot create union schema from empty inputs") + ); + } + + #[test] + fn test_union_single_input() -> Result<()> { + // Test that UnionExec::try_new returns the single input directly + let schema = create_test_schema()?; + let memory_exec: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + let memory_exec_clone = Arc::clone(&memory_exec); + let result = UnionExec::try_new(vec![memory_exec])?; + + // Check that the result is the same as the input (no UnionExec wrapper) + assert_eq!(result.schema(), schema); + // Verify it's the same execution plan + assert!(Arc::ptr_eq(&result, &memory_exec_clone)); + + Ok(()) + } + + #[test] + fn test_union_schema_multiple_inputs() -> Result<()> { + // Test that existing functionality with multiple inputs still works + let schema = create_test_schema()?; + let memory_exec1 = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + let memory_exec2 = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + + let union_plan = UnionExec::try_new(vec![memory_exec1, memory_exec2])?; + + // Downcast to verify it's a UnionExec + let union = union_plan + .downcast_ref::<UnionExec>() + .expect("Expected UnionExec"); + + // Check that schema is correct + assert_eq!(union.schema(), schema); + // Check that we have 2 inputs + assert_eq!(union.inputs().len(), 2); + + Ok(()) + } + + #[test] + fn test_union_schema_mismatch() { + // Test that UnionExec properly rejects inputs with different field counts + let schema = create_test_schema().unwrap(); + let schema2 = create_test_schema2().unwrap(); + let memory_exec1 = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None).unwrap()); + let memory_exec2 = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema2), None).unwrap()); + + let result = UnionExec::try_new(vec![memory_exec1, memory_exec2]); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains( + "UnionExec/InterleaveExec requires all inputs to have the same number of fields" + ) + ); + } + + fn make_hash_exec( + schema: &SchemaRef, + hash_cols: Vec<&str>, + buckets: usize, + ) -> Result<Arc<dyn ExecutionPlan>> { + let exprs = hash_cols + .iter() + .map(|c| col(c, schema)) + .collect::<Result<Vec<_>>>()?; + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Hash(exprs, buckets), + )?)) + } + + fn make_range_exec( + schema: &SchemaRef, + split_values: Vec<i32>, + sort_options: SortOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + let sort_expr = + PhysicalSortExpr::new(col(schema.field(0).name(), schema)?, sort_options); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let split_points = split_values + .into_iter() + .map(|v| SplitPoint::new(vec![ScalarValue::Int32(Some(v))])) + .collect(); + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?), + )?)) + } + + #[test] + fn test_can_interleave_matrix() -> Result<()> { + let name_column = "name"; + let age_column = "age"; + let schema = Arc::new(Schema::new(vec![ + Field::new(name_column, DataType::Int32, true), + Field::new(age_column, DataType::Int32, true), + ])); + + let ascending = SortOptions { + descending: false, + nulls_first: false, + }; + struct Case { + inputs: Vec<Arc<dyn ExecutionPlan>>, + expected: bool, + label: &'static str, + } + + let cases = vec![ + // compatible + Case { + label: "matching hash on single column", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + Case { + label: "matching hash on multiple columns", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + ], + }, + Case { + label: "matching range same splits and order", + expected: true, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 20], ascending)?, + ], + }, + // incompatible + Case { + label: "subset range partition", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 15], ascending)?, + ], + }, + Case { + label: "range different split points", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 30], ascending)?, + ], + }, + Case { + label: "mixed range and hash", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + ]; + + for case in cases { + assert_eq!( + can_interleave(case.inputs.iter()), + case.expected, + "{}", + case.label + ); + } + Ok(()) + } + + #[test] + fn test_union_cardinality_effect() -> Result<()> { + let schema = create_test_schema()?; + let input1: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + let input2: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + + let union = UnionExec::try_new(vec![input1, input2])?; + let union = union + .downcast_ref::<UnionExec>() + .expect("expected UnionExec for multiple inputs"); + + assert!(matches!( + union.cardinality_effect(), + CardinalityEffect::GreaterEqual + )); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/unnest.rs b/native/vendor/datafusion-physical-plan/src/unnest.rs new file mode 100644 index 0000000..6dfc2a0 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/unnest.rs
@@ -0,0 +1,2408 @@ +// 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. + +//! Define a plan for unnesting values in columns that contain a list type. + +use std::cmp::{self, Ordering}; +use std::sync::Arc; +use std::task::{Poll, ready}; + +use super::metrics::{ + self, BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + MetricsSet, RecordOutput, SplitMetrics, +}; +use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; +use crate::stream::{BatchSplitStream, EmptyRecordBatchStream}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + validate_child_count, +}; + +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeListArray, Int64Array, + LargeListArray, LargeListViewArray, ListArray, ListViewArray, PrimitiveArray, Scalar, + StructArray, new_null_array, +}; +use arrow::compute::kernels::length::length; +use arrow::compute::kernels::zip::zip; +use arrow::compute::{cast, is_not_null, kernels, sum}; +use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use arrow_ord::cmp::lt; +use async_trait::async_trait; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{ + Constraints, HashMap, HashSet, Result, UnnestOptions, exec_datafusion_err, exec_err, + internal_err, +}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr::expressions::Column; +use futures::{Stream, StreamExt}; +use log::trace; + +/// Unnest the given columns (either with type struct or list) +/// For list unnesting, each row is vertically transformed into multiple rows +/// For struct unnesting, each column is horizontally transformed into multiple columns, +/// Thus the original RecordBatch with dimension (n x m) may have new dimension (n' x m') +/// +/// See [`UnnestOptions`] for more details and an example. +#[derive(Debug, Clone)] +pub struct UnnestExec { + /// Input execution plan + input: Arc<dyn ExecutionPlan>, + /// The schema once the unnest is applied + schema: SchemaRef, + /// Indices of the list-typed columns in the input schema + list_column_indices: Vec<ListUnnest>, + /// Indices of the struct-typed columns in the input schema + struct_column_indices: Vec<usize>, + /// Options + options: UnnestOptions, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl UnnestExec { + /// Create a new [UnnestExec]. + pub fn new( + input: Arc<dyn ExecutionPlan>, + list_column_indices: Vec<ListUnnest>, + struct_column_indices: Vec<usize>, + schema: SchemaRef, + options: UnnestOptions, + ) -> Result<Self> { + let cache = Self::compute_properties( + &input, + &list_column_indices, + &struct_column_indices, + &schema, + )?; + + Ok(UnnestExec { + input, + schema, + list_column_indices, + struct_column_indices, + options, + metrics: Default::default(), + cache: Arc::new(cache), + }) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + list_column_indices: &[ListUnnest], + struct_column_indices: &[usize], + schema: &SchemaRef, + ) -> Result<PlanProperties> { + // Find out which indices are not unnested, such that they can be copied over from the input plan + let input_schema = input.schema(); + let mut unnested_indices = BooleanBufferBuilder::new(input_schema.fields().len()); + unnested_indices.append_n(input_schema.fields().len(), false); + for list_unnest in list_column_indices { + unnested_indices.set_bit(list_unnest.index_in_input_schema, true); + } + for struct_unnest in struct_column_indices { + unnested_indices.set_bit(*struct_unnest, true) + } + let unnested_indices = unnested_indices.finish(); + let non_unnested_indices: Vec<usize> = (0..input_schema.fields().len()) + .filter(|idx| !unnested_indices.value(*idx)) + .collect(); + + // Manually build projection mapping from non-unnested input columns to their positions in the output + let input_schema = input.schema(); + let projection_mapping: ProjectionMapping = non_unnested_indices + .iter() + .map(|&input_idx| { + // Find what index the input column has in the output schema + let input_field = input_schema.field(input_idx); + let output_idx = schema + .fields() + .iter() + .position(|output_field| output_field.name() == input_field.name()) + .ok_or_else(|| { + exec_datafusion_err!( + "Non-unnested column '{}' must exist in output schema", + input_field.name() + ) + })?; + + let input_col = Arc::new(Column::new(input_field.name(), input_idx)) + as Arc<dyn PhysicalExpr>; + let target_col = Arc::new(Column::new(input_field.name(), output_idx)) + as Arc<dyn PhysicalExpr>; + // Use From<Vec<(Arc<dyn PhysicalExpr>, usize)>> for ProjectionTargets + let targets = vec![(target_col, output_idx)].into(); + Ok((input_col, targets)) + }) + .collect::<Result<ProjectionMapping>>()?; + + // Create the unnest's equivalence properties by copying the input plan's equivalence properties + // for the unaffected columns. Except for the constraints, which are removed entirely because + // the unnest operation invalidates any global uniqueness or primary-key constraints. + let input_eq_properties = input.equivalence_properties(); + let eq_properties = input_eq_properties + .project(&projection_mapping, Arc::clone(schema)) + .with_constraints(Constraints::default()); + + // Output partitioning must use the projection mapping + let output_partitioning = input + .output_partitioning() + .project(&projection_mapping, &eq_properties); + + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + input.pipeline_behavior(), + input.boundedness(), + )) + } + + /// Input execution plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Indices of the list-typed columns in the input schema + pub fn list_column_indices(&self) -> &[ListUnnest] { + &self.list_column_indices + } + + /// Indices of the struct-typed columns in the input schema + pub fn struct_column_indices(&self) -> &[usize] { + &self.struct_column_indices + } + + pub fn options(&self) -> &UnnestOptions { + &self.options + } +} + +impl DisplayAs for UnnestExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "UnnestExec") + } + DisplayFormatType::TreeRender => { + write!(f, "") + } + } + } +} + +impl ExecutionPlan for UnnestExec { + fn name(&self) -> &'static str { + "UnnestExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(UnnestExec::new( + children.swap_remove(0), + self.list_column_indices.clone(), + self.struct_column_indices.clone(), + Arc::clone(&self.schema), + self.options.clone(), + )?)), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let batch_size = context.session_config().batch_size(); + let input = self.input.execute(partition, context)?; + let metrics = UnnestMetrics::new(partition, &self.metrics); + + let stream = Box::pin(UnnestStream { + input, + schema: Arc::clone(&self.schema), + list_type_columns: self.list_column_indices.clone(), + struct_column_indices: self.struct_column_indices.iter().copied().collect(), + options: self.options.clone(), + metrics, + batch_size, + pending_input: None, + }); + + // Chunking the input bounds each build to roughly `batch_size` rows, but two cases + // can still produce an oversized batch (see `predict_output_lens`), so the output + // goes through the shared splitter to make the bound unconditional. + Ok(Box::pin(BatchSplitStream::new( + stream, + batch_size, + SplitMetrics::new(&self.metrics, partition), + ))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `UnnestExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + input, + schema, + list_column_indices, + struct_column_indices, + options, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + // Derived at construction by `UnnestExec::compute_properties`. + cache: _, + } = self; + + let input = ctx.encode_child(input)?; + let schema = schema.as_ref().try_into()?; + let list_type_columns = list_column_indices + .iter() + .map(|column| protobuf::ListUnnest { + index_in_input_schema: column.index_in_input_schema as _, + depth: column.depth as _, + }) + .collect(); + let struct_type_columns = struct_column_indices + .iter() + .map(|index| *index as _) + .collect(); + let null_handling = { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + match options.null_handling { + NullHandling::Preserve => ProtoNullHandling::Preserve, + NullHandling::Drop => ProtoNullHandling::Drop, + NullHandling::PreserveAndExpandEmpty => { + ProtoNullHandling::PreserveAndExpandEmpty + } + } + } as i32; + let options = protobuf::UnnestOptions { + null_handling, + recursions: options + .recursions + .iter() + .map(|recursion| protobuf::RecursionUnnestOption { + input_column: Some((&recursion.input_column).into()), + output_column: Some((&recursion.output_column).into()), + depth: recursion.depth as _, + }) + .collect(), + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Unnest(Box::new( + protobuf::UnnestExecNode { + input: Some(Box::new(input)), + schema: Some(schema), + list_type_columns, + struct_type_columns, + options: Some(options), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnnestExec { + /// Reconstruct an [`UnnestExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use datafusion_proto_models::protobuf; + + let unnest = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Unnest, + "UnnestExec", + ); + // Exhaustive destructure: a new field on `UnnestExecNode` is a compile + // error here rather than a silently ignored wire field. + let protobuf::UnnestExecNode { + input, + schema, + list_type_columns, + struct_type_columns, + options, + } = unnest.as_ref(); + + let input = ctx.decode_required_child(input.as_deref(), "UnnestExec", "input")?; + let schema: Schema = schema + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "UnnestExec is missing required field 'schema'" + ) + })? + .try_into()?; + let list_column_indices = list_type_columns + .iter() + .map(|column| ListUnnest { + index_in_input_schema: column.index_in_input_schema as _, + depth: column.depth as _, + }) + .collect(); + let struct_column_indices = struct_type_columns + .iter() + .map(|index| *index as _) + .collect(); + let options = options.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "UnnestExec is missing required field 'options'" + ) + })?; + let null_handling = { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + match ProtoNullHandling::try_from(options.null_handling) { + Ok(ProtoNullHandling::Preserve) => NullHandling::Preserve, + Ok(ProtoNullHandling::Drop) => NullHandling::Drop, + Ok(ProtoNullHandling::PreserveAndExpandEmpty) => { + NullHandling::PreserveAndExpandEmpty + } + // Unknown enum values fall back to the default (Preserve), + // matching DataFusion's historical behavior. + Err(_) => NullHandling::Preserve, + } + }; + let options = UnnestOptions { + null_handling, + recursions: options + .recursions + .iter() + .map(|recursion| datafusion_common::RecursionUnnestOption { + input_column: recursion.input_column.as_ref().unwrap().into(), + output_column: recursion.output_column.as_ref().unwrap().into(), + depth: recursion.depth as _, + }) + .collect(), + }; + + Ok(Arc::new(UnnestExec::new( + input, + list_column_indices, + struct_column_indices, + Arc::new(schema), + options, + )?)) + } +} + +#[derive(Clone, Debug)] +struct UnnestMetrics { + /// Execution metrics + baseline_metrics: BaselineMetrics, + /// Number of batches consumed + input_batches: metrics::Count, + /// Number of rows consumed + input_rows: metrics::Count, +} + +impl UnnestMetrics { + fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { + let input_batches = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("input_batches", partition); + + let input_rows = MetricBuilder::new(metrics) + .with_category(MetricCategory::Rows) + .counter("input_rows", partition); + + Self { + baseline_metrics: BaselineMetrics::new(metrics, partition), + input_batches, + input_rows, + } + } +} + +/// A stream that issues [RecordBatch]es with unnested column data. +struct UnnestStream { + /// Input stream + input: SendableRecordBatchStream, + /// Unnested schema + schema: Arc<Schema>, + /// represents all unnest operations to be applied to the input (input index, depth) + /// e.g unnest(col1),unnest(unnest(col1)) where col1 has index 1 in original input schema + /// then list_type_columns = [ListUnnest{1,1},ListUnnest{1,2}] + list_type_columns: Vec<ListUnnest>, + struct_column_indices: HashSet<usize>, + /// Options + options: UnnestOptions, + /// Metrics + metrics: UnnestMetrics, + /// Target number of rows per output batch, from `datafusion.execution.batch_size`. + batch_size: usize, + /// Rows of the current input batch that have not been unnested yet. Unnesting one + /// input batch can produce arbitrarily many output rows, so the input is consumed in + /// chunks small enough that each chunk's output stays near `batch_size`. + /// + /// Note the scope of the memory bound this buys: chunking removes the input batch size + /// from the peak, but not the length of an individual list. A single row whose list is + /// longer than `batch_size`, and recursive unnesting (where the expansion cannot be + /// predicted up front), both still materialize their full expansion in one build. + pending_input: Option<PendingInput>, +} + +/// An input batch being unnested incrementally, a chunk of rows at a time. +struct PendingInput { + /// The full input batch. Rows before `row_offset` have already been unnested. + batch: RecordBatch, + /// Index of the next input row to unnest. + row_offset: usize, + /// How many output rows each input row expands into, indexed by input row. + /// + /// `None` when the expansion cannot be predicted from the input alone, in which case + /// the whole remaining input is unnested in one call and only the output is split. + /// See [`UnnestStream::predict_output_lens`]. + output_lens: Option<PrimitiveArray<Int64Type>>, +} + +impl PendingInput { + fn remaining_rows(&self) -> usize { + self.batch.num_rows() - self.row_offset + } + + /// How many input rows to unnest next so the resulting batch holds at most + /// `batch_size` rows. + fn next_chunk_rows(&self, batch_size: usize) -> usize { + let Some(output_lens) = &self.output_lens else { + return self.remaining_rows(); + }; + + let lens = &output_lens.values()[self.row_offset..]; + let batch_size = batch_size as i64; + let mut output_rows = 0i64; + for (rows, len) in lens.iter().enumerate() { + // The first row is always taken, even if it alone overshoots `batch_size`: an + // input row is never split across builds, so this is what guarantees progress. + // An oversized build is sliced down by `BatchSplitStream` on the way out. + if rows > 0 && output_rows + len > batch_size { + return rows; + } + output_rows += len; + } + lens.len() + } + + /// The per-row output lengths covering the next `rows` input rows, so the unnesting + /// does not have to recompute what `predict_output_lens` already derived. + fn chunk_lengths(&self, rows: usize) -> Option<PrimitiveArray<Int64Type>> { + self.output_lens + .as_ref() + .map(|lens| lens.slice(self.row_offset, rows)) + } +} + +impl RecordBatchStream for UnnestStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[async_trait] +impl Stream for UnnestStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_impl(cx) + } +} + +impl UnnestStream { + /// Separate implementation function that unpins the [`UnnestStream`] so + /// that partial borrows work correctly + fn poll_next_impl( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + // Unnest the next chunk of the input batch already in hand. + if let Some(pending) = self.pending_input.as_mut() { + // `PendingInput` is only built from a non-empty batch and `next_chunk_rows` + // always consumes at least one row, so it is dropped the moment it drains. + debug_assert!(pending.remaining_rows() > 0); + + let rows = pending.next_chunk_rows(self.batch_size); + let chunk = pending.batch.slice(pending.row_offset, rows); + let chunk_lengths = pending.chunk_lengths(rows); + pending.row_offset += rows; + let drained = pending.remaining_rows() == 0; + + let timer = self.metrics.baseline_metrics.elapsed_compute().timer(); + let result = build_batch( + &chunk, + &self.schema, + &self.list_type_columns, + &self.struct_column_indices, + &self.options, + chunk_lengths.as_ref(), + ); + timer.done(); + + if drained { + self.pending_input = None; + } + + // A chunk can legitimately produce no rows at all, for example when every + // list in it is empty under `NullHandling::Drop`; `build_batch` signals + // that with `None` rather than an empty batch. + if let Some(batch) = result? { + debug_assert!(batch.num_rows() > 0); + (&batch).record_output(&self.metrics.baseline_metrics); + return Poll::Ready(Some(Ok(batch))); + } + continue; + } + + // Otherwise pull the next input batch. + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + self.metrics.input_batches.add(1); + self.metrics.input_rows.add(batch.num_rows()); + if batch.num_rows() > 0 { + let timer = + self.metrics.baseline_metrics.elapsed_compute().timer(); + let output_lens = self.predict_output_lens(&batch); + timer.done(); + self.pending_input = Some(PendingInput { + batch, + row_offset: 0, + output_lens: output_lens?, + }); + } + } + // If the stream is depleted or returned an error, log the finish message: + other => { + trace!( + "Processed {} probe-side input batches containing {} rows and \ + produced {} output batches containing {} rows in {}", + self.metrics.input_batches, + self.metrics.input_rows, + self.metrics.baseline_metrics.output_batches(), + self.metrics.baseline_metrics.output_rows(), + self.metrics.baseline_metrics.elapsed_compute(), + ); + + // In the non-error case, i.e., input is simply depleted: + if other.is_none() { + // Release the input pipeline's resources. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + return Poll::Ready(other); + } + } + } + } + + /// Compute how many output rows each input row of `batch` will expand into, so the + /// input can be chunked to keep each build bounded. + /// + /// Returns `None` when the count cannot be derived from the input alone, which is the + /// signal to unnest the whole batch in one call: + /// + /// * With no list columns, unnesting only widens structs and leaves the row count + /// alone, so the output is already bounded by the input batch size. + /// * With recursion (`depth > 1`), a row's expansion depends on the lengths of inner + /// lists that only exist after the outer levels have been unnested, so it cannot be + /// predicted up front. + fn predict_output_lens( + &self, + batch: &RecordBatch, + ) -> Result<Option<PrimitiveArray<Int64Type>>> { + if self.list_type_columns.is_empty() + || self + .list_type_columns + .iter() + .any(|unnest| unnest.depth != 1) + { + return Ok(None); + } + + let list_arrays: Vec<ArrayRef> = self + .list_type_columns + .iter() + .map(|unnest| Arc::clone(batch.column(unnest.index_in_input_schema))) + .collect(); + + // This is exactly the per-row length that `list_unnest_at_level` derives when it + // actually unnests, so the chunk boundaries are exact rather than estimated, and + // each chunk's slice of it is handed back to `build_batch` instead of recomputed. + let longest_length = find_longest_length(&list_arrays, &self.options)?; + Ok(Some(longest_length.as_primitive::<Int64Type>().clone())) + } +} + +/// Given a set of struct column indices to flatten +/// try converting the column in input into multiple subfield columns +/// For example +/// struct_col: [a: struct(item: int, name: string), b: int] +/// with a batch +/// {a: {item: 1, name: "a"}, b: 2}, +/// {a: {item: 3, name: "b"}, b: 4] +/// will be converted into +/// {a.item: 1, a.name: "a", b: 2}, +/// {a.item: 3, a.name: "b", b: 4} +fn flatten_struct_cols( + input_batch: &[Arc<dyn Array>], + schema: &SchemaRef, + struct_column_indices: &HashSet<usize>, +) -> Result<RecordBatch> { + // horizontal expansion because of struct unnest + let columns_expanded = input_batch + .iter() + .enumerate() + .map(|(idx, column_data)| match struct_column_indices.get(&idx) { + Some(_) => match column_data.data_type() { + DataType::Struct(_) => { + let struct_arr = + column_data.as_any().downcast_ref::<StructArray>().unwrap(); + Ok(struct_arr.columns().to_vec()) + } + data_type => internal_err!( + "expecting column {idx} from input plan to be a struct, got {data_type}" + ), + }, + None => Ok(vec![Arc::clone(column_data)]), + }) + .collect::<Result<Vec<_>>>()? + .into_iter() + .flatten() + .collect(); + Ok(RecordBatch::try_new(Arc::clone(schema), columns_expanded)?) +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] +pub struct ListUnnest { + pub index_in_input_schema: usize, + pub depth: usize, +} + +/// This function is used to execute the unnesting on multiple columns all at once, but +/// one level at a time, and is called n times, where n is the highest recursion level among +/// the unnest exprs in the query. +/// +/// For example giving the following query: +/// ```sql +/// select unnest(colA, max_depth:=3) as P1, unnest(colA,max_depth:=2) as P2, unnest(colB, max_depth:=1) as P3 from temp; +/// ``` +/// Then the total times this function being called is 3 +/// +/// It needs to be aware of which level the current unnesting is, because if there exists +/// multiple unnesting on the same column, but with different recursion levels, say +/// **unnest(colA, max_depth:=3)** and **unnest(colA, max_depth:=2)**, then the unnesting +/// of expr **unnest(colA, max_depth:=3)** will start at level 3, while unnesting for expr +/// **unnest(colA, max_depth:=2)** has to start at level 2 +/// +/// Set *colA* as a 3-dimension columns and *colB* as an array (1-dimension). As stated, +/// this function is called with the descending order of recursion depth +/// +/// Depth = 3 +/// - colA(3-dimension) unnest into temp column temp_P1(2_dimension) (unnesting of P1 starts +/// from this level) +/// - colA(3-dimension) having indices repeated by the unnesting operation above +/// - colB(1-dimension) having indices repeated by the unnesting operation above +/// +/// Depth = 2 +/// - temp_P1(2-dimension) unnest into temp column temp_P1(1-dimension) +/// - colA(3-dimension) unnest into temp column temp_P2(2-dimension) (unnesting of P2 starts +/// from this level) +/// - colB(1-dimension) having indices repeated by the unnesting operation above +/// +/// Depth = 1 +/// - temp_P1(1-dimension) unnest into P1 +/// - temp_P2(2-dimension) unnest into P2 +/// - colB(1-dimension) unnest into P3 (unnesting of P3 starts from this level) +/// +/// The returned array will has the same size as the input batch +/// and only contains original columns that are not being unnested. +fn list_unnest_at_level( + batch: &[ArrayRef], + list_type_unnests: &[ListUnnest], + temp_unnested_arrs: &mut HashMap<ListUnnest, ArrayRef>, + level_to_unnest: usize, + options: &UnnestOptions, + precomputed_lengths: Option<&PrimitiveArray<Int64Type>>, +) -> Result<Option<Vec<ArrayRef>>> { + // Extract unnestable columns at this level + let (arrs_to_unnest, list_unnest_specs): (Vec<Arc<dyn Array>>, Vec<_>) = + list_type_unnests + .iter() + .filter_map(|unnesting| { + if level_to_unnest == unnesting.depth { + return Some(( + Arc::clone(&batch[unnesting.index_in_input_schema]), + *unnesting, + )); + } + // This means the unnesting on this item has started at higher level + // and need to continue until depth reaches 1 + if level_to_unnest < unnesting.depth { + return Some(( + Arc::clone(temp_unnested_arrs.get(unnesting).unwrap()), + *unnesting, + )); + } + None + }) + .unzip(); + + // Filter out so that list_arrays only contain column with the highest depth + // at the same time, during iteration remove this depth so next time we don't have to unnest them again + // + // The caller may already have computed these lengths to decide how many input rows to + // feed us; reusing them avoids running the kernel chain twice over the same rows. + // Cloning is an `Arc` bump on the underlying buffer, not a copy. + let longest_length = match precomputed_lengths { + Some(lengths) => lengths.clone(), + None => find_longest_length(&arrs_to_unnest, options)? + .as_primitive::<Int64Type>() + .clone(), + }; + let unnested_length = &longest_length; + let total_length = if unnested_length.is_empty() { + 0 + } else { + sum(unnested_length).ok_or_else(|| { + exec_datafusion_err!("Failed to calculate the total unnested length") + })? as usize + }; + if total_length == 0 { + return Ok(None); + } + + // Unnest all the list arrays + let unnested_temp_arrays = + unnest_list_arrays(arrs_to_unnest.as_ref(), unnested_length, total_length)?; + + // Create the take indices array for other columns + let take_indices = create_take_indices(unnested_length, total_length); + unnested_temp_arrays + .into_iter() + .zip(list_unnest_specs.iter()) + .for_each(|(flatten_arr, unnesting)| { + temp_unnested_arrs.insert(*unnesting, flatten_arr); + }); + + let repeat_mask: Vec<bool> = batch + .iter() + .enumerate() + .map(|(i, _)| { + // Check if the column is needed in future levels (levels below the current one) + let needed_in_future_levels = list_type_unnests.iter().any(|unnesting| { + unnesting.index_in_input_schema == i && unnesting.depth < level_to_unnest + }); + + // Check if the column is involved in unnesting at any level + let is_involved_in_unnesting = list_type_unnests + .iter() + .any(|unnesting| unnesting.index_in_input_schema == i); + + // Repeat columns needed in future levels or not unnested. + needed_in_future_levels || !is_involved_in_unnesting + }) + .collect(); + + // Dimension of arrays in batch is untouched, but the values are repeated + // as the side effect of unnesting + let ret = repeat_arrs_from_indices(batch, &take_indices, &repeat_mask)?; + + Ok(Some(ret)) +} +struct UnnestingResult { + arr: ArrayRef, + depth: usize, +} + +/// For each row in a `RecordBatch`, some list/struct columns need to be unnested. +/// - For list columns: We will expand the values in each list into multiple rows, +/// taking the longest length among these lists, and shorter lists are padded with NULLs. +/// - For struct columns: We will expand the struct columns into multiple subfield columns. +/// +/// For columns that don't need to be unnested, repeat their values until reaching the longest length. +/// +/// Note: unnest has a big difference in behavior between Postgres and DuckDB +/// +/// Take this example +/// +/// 1. Postgres +/// ```ignored +/// create table temp ( +/// i integer[][][], j integer[] +/// ) +/// insert into temp values ('{{{1,2},{3,4}},{{5,6},{7,8}}}', '{1,2}'); +/// select unnest(i), unnest(j) from temp; +/// ``` +/// +/// Result +/// ```text +/// 1 1 +/// 2 2 +/// 3 +/// 4 +/// 5 +/// 6 +/// 7 +/// 8 +/// ``` +/// 2. DuckDB +/// ```ignore +/// create table temp (i integer[][][], j integer[]); +/// insert into temp values ([[[1,2],[3,4]],[[5,6],[7,8]]], [1,2]); +/// select unnest(i,recursive:=true), unnest(j,recursive:=true) from temp; +/// ``` +/// Result: +/// ```text +/// +/// ┌────────────────────────────────────────────────┬────────────────────────────────────────────────┐ +/// │ unnest(i, "recursive" := CAST('t' AS BOOLEAN)) │ unnest(j, "recursive" := CAST('t' AS BOOLEAN)) │ +/// │ int32 │ int32 │ +/// ├────────────────────────────────────────────────┼────────────────────────────────────────────────┤ +/// │ 1 │ 1 │ +/// │ 2 │ 2 │ +/// │ 3 │ 1 │ +/// │ 4 │ 2 │ +/// │ 5 │ 1 │ +/// │ 6 │ 2 │ +/// │ 7 │ 1 │ +/// │ 8 │ 2 │ +/// └────────────────────────────────────────────────┴────────────────────────────────────────────────┘ +/// ``` +/// +/// The following implementation refer to DuckDB's implementation +fn build_batch( + batch: &RecordBatch, + schema: &SchemaRef, + list_type_columns: &[ListUnnest], + struct_column_indices: &HashSet<usize>, + options: &UnnestOptions, + precomputed_lengths: Option<&PrimitiveArray<Int64Type>>, +) -> Result<Option<RecordBatch>> { + let transformed = match list_type_columns.len() { + 0 => flatten_struct_cols(batch.columns(), schema, struct_column_indices), + _ => { + let mut temp_unnested_result = HashMap::new(); + let max_recursion = list_type_columns + .iter() + .fold(0, |highest_depth, ListUnnest { depth, .. }| { + cmp::max(highest_depth, *depth) + }); + + // This arr always has the same column count with the input batch + let mut flatten_arrs = vec![]; + + // Original batch has the same columns + // All unnesting results are written to temp_batch + for depth in (1..=max_recursion).rev() { + let input = match depth == max_recursion { + true => batch.columns(), + false => &flatten_arrs, + }; + // Only sound for a single non-recursive level: with recursion the deeper + // levels' lengths depend on arrays that do not exist yet, which is also why + // the caller does not predict lengths in that case. + let level_lengths = if max_recursion == 1 { + precomputed_lengths + } else { + None + }; + let Some(temp_result) = list_unnest_at_level( + input, + list_type_columns, + &mut temp_unnested_result, + depth, + options, + level_lengths, + )? + else { + return Ok(None); + }; + flatten_arrs = temp_result; + } + let unnested_array_map: HashMap<usize, Vec<UnnestingResult>> = + temp_unnested_result.into_iter().fold( + HashMap::new(), + |mut acc, + ( + ListUnnest { + index_in_input_schema, + depth, + }, + flattened_array, + )| { + acc.entry(index_in_input_schema).or_default().push( + UnnestingResult { + arr: flattened_array, + depth, + }, + ); + acc + }, + ); + let output_order: HashMap<ListUnnest, usize> = list_type_columns + .iter() + .enumerate() + .map(|(order, unnest_def)| (*unnest_def, order)) + .collect(); + + // One original column may be unnested multiple times into separate columns + let mut multi_unnested_per_original_index = unnested_array_map + .into_iter() + .map( + // Each item in unnested_columns is the result of unnesting the same input column + // we need to sort them to conform with the original expression order + // e.g unnest(unnest(col)) must goes before unnest(col) + |(original_index, mut unnested_columns)| { + unnested_columns.sort_by( + |UnnestingResult { depth: depth1, .. }, + UnnestingResult { depth: depth2, .. }| + -> Ordering { + output_order + .get(&ListUnnest { + depth: *depth1, + index_in_input_schema: original_index, + }) + .unwrap() + .cmp( + output_order + .get(&ListUnnest { + depth: *depth2, + index_in_input_schema: original_index, + }) + .unwrap(), + ) + }, + ); + ( + original_index, + unnested_columns + .into_iter() + .map(|result| result.arr) + .collect::<Vec<_>>(), + ) + }, + ) + .collect::<HashMap<_, _>>(); + + let ret = flatten_arrs + .into_iter() + .enumerate() + .flat_map(|(col_idx, arr)| { + // Convert original column into its unnested version(s) + // Plural because one column can be unnested with different recursion level + // and into separate output columns + match multi_unnested_per_original_index.remove(&col_idx) { + Some(unnested_arrays) => unnested_arrays, + None => vec![arr], + } + }) + .collect::<Vec<_>>(); + + flatten_struct_cols(&ret, schema, struct_column_indices) + } + }?; + Ok(Some(transformed)) +} + +/// Find the longest list length among the given list arrays for each row. +/// +/// For example if we have the following two list arrays: +/// +/// ```ignore +/// l1: [1, 2, 3], null, [], [3] +/// l2: [4,5], [], null, [6, 7] +/// ``` +/// +/// With [`datafusion_common::NullHandling::Drop`], the longest length array will be: +/// +/// ```ignore +/// longest_length: [3, 0, 0, 2] +/// ``` +/// +/// With [`datafusion_common::NullHandling::Preserve`] (the default), the longest length array +/// will be: +/// +/// ```ignore +/// longest_length: [3, 1, 1, 2] +/// ``` +/// +/// With [`datafusion_common::NullHandling::PreserveAndExpandEmpty`], empty input lists are +/// also bumped to length 1 so they produce a single `NULL` output row: +/// +/// ```ignore +/// longest_length: [3, 1, 1, 2] +/// ``` +fn find_longest_length( + list_arrays: &[ArrayRef], + options: &UnnestOptions, +) -> Result<ArrayRef> { + // The length to substitute for a NULL input list. + let null_length = if options.preserve_nulls() { + Scalar::new(Int64Array::from_value(1, 1)) + } else { + Scalar::new(Int64Array::from_value(0, 1)) + }; + let expand_empty = options.expand_empty_as_null(); + // Reused scalars for the empty-list rewrite when expand_empty is set. + let zero = Scalar::new(Int64Array::from_value(0, 1)); + let one = Scalar::new(Int64Array::from_value(1, 1)); + let list_lengths: Vec<ArrayRef> = list_arrays + .iter() + .map(|list_array| { + let mut length_array = length(list_array)?; + // Make sure length arrays have the same type. Int64 is the most general one. + length_array = cast(&length_array, &DataType::Int64)?; + length_array = + zip(&is_not_null(&length_array)?, &length_array, &null_length)?; + if expand_empty { + // Bump empty lists (length 0) to length 1 so they + // produce a single output row padded with NULL. + let is_zero = arrow_ord::cmp::eq(&length_array, &zero)?; + length_array = zip(&is_zero, &one, &length_array)?; + } + Ok(length_array) + }) + .collect::<Result<_>>()?; + + let longest_length = list_lengths.iter().skip(1).try_fold( + Arc::clone(&list_lengths[0]), + |longest, current| { + let is_lt = lt(&longest, ¤t)?; + zip(&is_lt, ¤t, &longest) + }, + )?; + Ok(longest_length) +} + +/// Trait defining common methods used for unnesting, implemented by list array types. +trait ListArrayType: Array { + /// Returns a reference to the values of this list. + fn values(&self) -> &ArrayRef; + + /// Returns the start and end offset of the values for the given row. + fn value_offsets(&self, row: usize) -> (i64, i64); +} + +impl ListArrayType for ListArray { + fn values(&self) -> &ArrayRef { + self.values() + } + + fn value_offsets(&self, row: usize) -> (i64, i64) { + let offsets = self.value_offsets(); + (offsets[row].into(), offsets[row + 1].into()) + } +} + +impl ListArrayType for LargeListArray { + fn values(&self) -> &ArrayRef { + self.values() + } + + fn value_offsets(&self, row: usize) -> (i64, i64) { + let offsets = self.value_offsets(); + (offsets[row], offsets[row + 1]) + } +} + +impl ListArrayType for FixedSizeListArray { + fn values(&self) -> &ArrayRef { + self.values() + } + + fn value_offsets(&self, row: usize) -> (i64, i64) { + let start = self.value_offset(row) as i64; + (start, start + self.value_length() as i64) + } +} + +impl ListArrayType for ListViewArray { + fn values(&self) -> &ArrayRef { + self.values() + } + + fn value_offsets(&self, row: usize) -> (i64, i64) { + let offset = self.value_offsets()[row] as i64; + let size = self.value_sizes()[row] as i64; + (offset, offset + size) + } +} + +impl ListArrayType for LargeListViewArray { + fn values(&self) -> &ArrayRef { + self.values() + } + + fn value_offsets(&self, row: usize) -> (i64, i64) { + let offset = self.value_offsets()[row]; + let size = self.value_sizes()[row]; + (offset, offset + size) + } +} + +/// Unnest multiple list arrays according to the length array. +fn unnest_list_arrays( + list_arrays: &[ArrayRef], + length_array: &PrimitiveArray<Int64Type>, + capacity: usize, +) -> Result<Vec<ArrayRef>> { + let typed_arrays = list_arrays + .iter() + .map(|list_array| match list_array.data_type() { + DataType::List(_) => Ok(list_array.as_list::<i32>() as &dyn ListArrayType), + DataType::LargeList(_) => { + Ok(list_array.as_list::<i64>() as &dyn ListArrayType) + } + DataType::FixedSizeList(_, _) => { + Ok(list_array.as_fixed_size_list() as &dyn ListArrayType) + } + DataType::ListView(_) => { + Ok(list_array.as_list_view::<i32>() as &dyn ListArrayType) + } + DataType::LargeListView(_) => { + Ok(list_array.as_list_view::<i64>() as &dyn ListArrayType) + } + other => exec_err!("Invalid unnest datatype {other }"), + }) + .collect::<Result<Vec<_>>>()?; + + typed_arrays + .iter() + .map(|list_array| unnest_list_array(*list_array, length_array, capacity)) + .collect::<Result<_>>() +} + +/// Unnest a list array according the target length array. +/// +/// Consider a list array like this: +/// +/// ```ignore +/// [1], [2, 3, 4], null, [5], [], +/// ``` +/// +/// and the length array is: +/// +/// ```ignore +/// [2, 3, 2, 1, 2] +/// ``` +/// +/// If the length of a certain list is less than the target length, pad with NULLs. +/// So the unnested array will look like this: +/// +/// ```ignore +/// [1, null, 2, 3, 4, null, null, 5, null, null] +/// ``` +fn unnest_list_array( + list_array: &dyn ListArrayType, + length_array: &PrimitiveArray<Int64Type>, + capacity: usize, +) -> Result<ArrayRef> { + let values = list_array.values(); + let mut take_indices_builder = PrimitiveArray::<Int64Type>::builder(capacity); + for row in 0..list_array.len() { + let mut value_length = 0; + if !list_array.is_null(row) { + let (start, end) = list_array.value_offsets(row); + value_length = end - start; + for i in start..end { + take_indices_builder.append_value(i) + } + } + let target_length = length_array.value(row); + debug_assert!( + value_length <= target_length, + "value length is beyond the longest length" + ); + // Pad with NULL values + for _ in value_length..target_length { + take_indices_builder.append_null(); + } + } + Ok(kernels::take::take( + &values, + &take_indices_builder.finish(), + None, + )?) +} + +/// Creates take indices that will be used to expand all columns except for the list type +/// [`columns`](UnnestExec::list_column_indices) that is being unnested. +/// Every column value needs to be repeated multiple times according to the length array. +/// +/// If the length array looks like this: +/// +/// ```ignore +/// [2, 3, 1] +/// ``` +/// Then [`create_take_indices`] will return an array like this +/// +/// ```ignore +/// [0, 0, 1, 1, 1, 2] +/// ``` +fn create_take_indices( + length_array: &PrimitiveArray<Int64Type>, + capacity: usize, +) -> PrimitiveArray<Int64Type> { + // `find_longest_length()` guarantees this. + debug_assert!( + length_array.null_count() == 0, + "length array should not contain nulls" + ); + let mut builder = PrimitiveArray::<Int64Type>::builder(capacity); + for (index, repeat) in length_array.iter().enumerate() { + // The length array should not contain nulls, so unwrap is safe + let repeat = repeat.unwrap(); + (0..repeat).for_each(|_| builder.append_value(index as i64)); + } + builder.finish() +} + +/// Create a batch of arrays based on an input `batch` and a `indices` array. +/// The `indices` array is used by the take kernel to repeat values in the arrays +/// that are marked with `true` in the `repeat_mask`. Arrays marked with `false` +/// in the `repeat_mask` will be replaced with arrays filled with nulls of the +/// appropriate length. +/// +/// For example if we have the following batch: +/// +/// ```ignore +/// c1: [1], null, [2, 3, 4], null, [5, 6] +/// c2: 'a', 'b', 'c', null, 'd' +/// ``` +/// +/// then the `unnested_list_arrays` contains the unnest column that will replace `c1` in +/// the final batch if `preserve_nulls` is true: +/// +/// ```ignore +/// c1: 1, null, 2, 3, 4, null, 5, 6 +/// ``` +/// +/// And the `indices` array contains the indices that are used by `take` kernel to +/// repeat the values in `c2`: +/// +/// ```ignore +/// 0, 1, 2, 2, 2, 3, 4, 4 +/// ``` +/// +/// so that the final batch will look like: +/// +/// ```ignore +/// c1: 1, null, 2, 3, 4, null, 5, 6 +/// c2: 'a', 'b', 'c', 'c', 'c', null, 'd', 'd' +/// ``` +/// +/// The `repeat_mask` determines whether an array's values are repeated or replaced with nulls. +/// For example, if the `repeat_mask` is: +/// +/// ```ignore +/// [true, false] +/// ``` +/// +/// The final batch will look like: +/// +/// ```ignore +/// c1: 1, null, 2, 3, 4, null, 5, 6 // Repeated using `indices` +/// c2: null, null, null, null, null, null, null, null // Replaced with nulls +fn repeat_arrs_from_indices( + batch: &[ArrayRef], + indices: &PrimitiveArray<Int64Type>, + repeat_mask: &[bool], +) -> Result<Vec<Arc<dyn Array>>> { + batch + .iter() + .zip(repeat_mask.iter()) + .map(|(arr, &repeat)| { + if repeat { + Ok(kernels::take::take(arr, indices, None)?) + } else { + Ok(new_null_array(arr.data_type(), arr.len())) + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + GenericListArray, Int32Array, NullBufferBuilder, OffsetSizeTrait, StringArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::datatypes::{Field, Int32Type}; + use datafusion_common::NullHandling; + use datafusion_common::test_util::batches_to_string; + use insta::assert_snapshot; + + // Create a GenericListArray with the following list values: + // [A, B, C], [], NULL, [D], NULL, [NULL, F] + fn make_generic_array<OffsetSize>() -> GenericListArray<OffsetSize> + where + OffsetSize: OffsetSizeTrait, + { + let mut values = vec![]; + let mut offsets: Vec<OffsetSize> = vec![OffsetSize::zero()]; + let mut valid = NullBufferBuilder::new(6); + + // [A, B, C] + values.extend_from_slice(&[Some("A"), Some("B"), Some("C")]); + offsets.push(OffsetSize::from_usize(values.len()).unwrap()); + valid.append_non_null(); + + // [] + offsets.push(OffsetSize::from_usize(values.len()).unwrap()); + valid.append_non_null(); + + // NULL with non-zero value length + // Issue https://github.com/apache/datafusion/issues/9932 + values.push(Some("?")); + offsets.push(OffsetSize::from_usize(values.len()).unwrap()); + valid.append_null(); + + // [D] + values.push(Some("D")); + offsets.push(OffsetSize::from_usize(values.len()).unwrap()); + valid.append_non_null(); + + // Another NULL with zero value length + offsets.push(OffsetSize::from_usize(values.len()).unwrap()); + valid.append_null(); + + // [NULL, F] + values.extend_from_slice(&[None, Some("F")]); + offsets.push(OffsetSize::from_usize(values.len()).unwrap()); + valid.append_non_null(); + + let field = Arc::new(Field::new_list_field(DataType::Utf8, true)); + GenericListArray::<OffsetSize>::new( + field, + OffsetBuffer::new(offsets.into()), + Arc::new(StringArray::from(values)), + valid.finish(), + ) + } + + // Create a FixedSizeListArray with the following list values: + // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] + fn make_fixed_list() -> FixedSizeListArray { + let values = Arc::new(StringArray::from_iter([ + Some("A"), + Some("B"), + None, + None, + Some("C"), + Some("D"), + None, + None, + None, + Some("F"), + None, + None, + ])); + let field = Arc::new(Field::new_list_field(DataType::Utf8, true)); + let valid = NullBuffer::from(vec![true, false, true, false, true, true]); + FixedSizeListArray::new(field, 2, values, Some(valid)) + } + + fn verify_unnest_list_array( + list_array: &dyn ListArrayType, + lengths: Vec<i64>, + expected: Vec<Option<&str>>, + ) -> Result<()> { + let length_array = Int64Array::from(lengths); + let unnested_array = unnest_list_array(list_array, &length_array, 3 * 6)?; + let strs = unnested_array.as_string::<i32>().iter().collect::<Vec<_>>(); + assert_eq!(strs, expected); + Ok(()) + } + + #[test] + fn test_build_batch_list_arr_recursive() -> Result<()> { + // col1 | col2 + // [[1,2,3],null,[4,5]] | ['a','b'] + // [[7,8,9,10], null, [11,12,13]] | ['c','d'] + // null | ['e'] + let list_arr1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + + let list_arr1_ref = Arc::new(list_arr1) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_null(); + // list<list<int32>> + let col1_field = Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + ))), + true, + ); + let col1 = ListArray::new( + Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr1_ref, + nulls.finish(), + ); + + let list_arr2 = StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ]); + + let offsets = OffsetBuffer::from_lengths([2, 2, 1]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_n_non_nulls(3); + let col2_field = Field::new( + "col2", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ); + let col2 = GenericListArray::<i32>::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(list_arr2), + nulls.finish(), + ); + // convert col1 and col2 to a record batch + let schema = Arc::new(Schema::new(vec![col1_field, col2_field])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new( + "col1_unnest_placeholder_depth_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new("col1_unnest_placeholder_depth_2", DataType::Int32, true), + Field::new("col2_unnest_placeholder_depth_1", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(col1) as ArrayRef, Arc::new(col2) as ArrayRef], + ) + .unwrap(); + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 0, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + let ret = build_batch( + &batch, + &out_schema, + list_type_columns.as_ref(), + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::Preserve, + recursions: vec![], + }, + None, + )? + .unwrap(); + + assert_snapshot!(batches_to_string(&[ret]), + @r" + +---------------------------------+---------------------------------+---------------------------------+ + | col1_unnest_placeholder_depth_1 | col1_unnest_placeholder_depth_2 | col2_unnest_placeholder_depth_1 | + +---------------------------------+---------------------------------+---------------------------------+ + | [1, 2, 3] | 1 | a | + | | 2 | b | + | [4, 5] | 3 | | + | [1, 2, 3] | | a | + | | | b | + | [4, 5] | | | + | [1, 2, 3] | 4 | a | + | | 5 | b | + | [4, 5] | | | + | [7, 8, 9, 10] | 7 | c | + | | 8 | d | + | [11, 12, 13] | 9 | | + | | 10 | | + | [7, 8, 9, 10] | | c | + | | | d | + | [11, 12, 13] | | | + | [7, 8, 9, 10] | 11 | c | + | | 12 | d | + | [11, 12, 13] | 13 | | + | | | e | + +---------------------------------+---------------------------------+---------------------------------+ + "); + Ok(()) + } + + #[test] + fn test_build_batch_preserve_and_expand_empty() -> Result<()> { + // c1: [A, B, C], [], NULL, [D], NULL, [NULL, F] c2: 1, 2, 3, 4, 5, 6 + // Expected for `NullHandling::PreserveAndExpandEmpty`: + // [A, B, C] -> three rows with c2 = 1, 1, 1 + // [] -> one row with c2 = 2 and unnested value NULL + // NULL -> one row with c2 = 3 and unnested value NULL + // [D] -> one row with c2 = 4 + // NULL -> one row with c2 = 5 and unnested value NULL + // [NULL, F] -> two rows with c2 = 6, 6 + let list_array = Arc::new(make_generic_array::<i32>()) as ArrayRef; + let other = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "c1", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("c2", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("c1_unnested", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![Arc::clone(&list_array), Arc::clone(&other)], + )?; + let list_type_columns = vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + None, + )? + .unwrap(); + + assert_snapshot!(batches_to_string(&[ret]), + @r" + +-------------+----+ + | c1_unnested | c2 | + +-------------+----+ + | A | 1 | + | B | 1 | + | C | 1 | + | | 2 | + | | 3 | + | D | 4 | + | | 5 | + | | 6 | + | F | 6 | + +-------------+----+ + "); + Ok(()) + } + + // PreserveAndExpandEmpty must work for LargeListArray (i64 offsets) too, + // not just the i32-offset ListArray exercised above. + #[test] + fn test_build_batch_preserve_and_expand_empty_largelist() -> Result<()> { + let list_array = Arc::new(make_generic_array::<i64>()) as ArrayRef; + let other = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "c1", + DataType::LargeList(Arc::new(Field::new_list_field( + DataType::Utf8, + true, + ))), + true, + ), + Field::new("c2", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("c1_unnested", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![Arc::clone(&list_array), Arc::clone(&other)], + )?; + let list_type_columns = vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + None, + )? + .unwrap(); + + // Same expected shape as the ListArray case — exercises the LargeList + // code path in unnest_list_array. + assert_snapshot!(batches_to_string(&[ret]), + @r" + +-------------+----+ + | c1_unnested | c2 | + +-------------+----+ + | A | 1 | + | B | 1 | + | C | 1 | + | | 2 | + | | 3 | + | D | 4 | + | | 5 | + | | 6 | + | F | 6 | + +-------------+----+ + "); + Ok(()) + } + + // When two list columns are unnested together, `find_longest_length` + // takes the per-row max. PreserveAndExpandEmpty must bump zeros to ones + // in each input column independently, then the row-wise max picks up + // the right value. + #[test] + fn test_build_batch_preserve_and_expand_empty_multi_column() -> Result<()> { + // col_a: [1, 2], [], NULL, [3] + // col_b: ['x'], ['y'],['z'], NULL + let col_a = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![]), + None, + Some(vec![Some(3)]), + ]); + let col_b = { + let mut b = + arrow::array::ListBuilder::new(arrow::array::StringBuilder::new()); + b.values().append_value("x"); + b.append(true); + b.values().append_value("y"); + b.append(true); + b.values().append_value("z"); + b.append(true); + b.append(false); + b.finish() + }; + let id = Arc::new(Int32Array::from(vec![10, 20, 30, 40])) as ArrayRef; + + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "a", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new( + "b", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("id", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("a_unnested", DataType::Int32, true), + Field::new("b_unnested", DataType::Utf8, true), + Field::new("id", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![ + Arc::new(col_a) as ArrayRef, + Arc::new(col_b) as ArrayRef, + Arc::clone(&id), + ], + )?; + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + None, + )? + .unwrap(); + + // Row 0: longest = max(len([1,2])=2, len(['x'])=1) = 2 → a=[1,2], b=['x',NULL] + // Row 1: a=[] bumped to len 1, b=['y'] len 1 → a=[NULL], b=['y'] + // Row 2: a=NULL bumped to len 1, b=['z'] len 1 → a=[NULL], b=['z'] + // Row 3: a=[3] len 1, b=NULL bumped to len 1 → a=[3], b=[NULL] + assert_snapshot!(batches_to_string(&[ret]), + @r" + +------------+------------+----+ + | a_unnested | b_unnested | id | + +------------+------------+----+ + | 1 | x | 10 | + | 2 | | 10 | + | | y | 20 | + | | z | 30 | + | 3 | | 40 | + +------------+------------+----+ + "); + Ok(()) + } + + // PreserveAndExpandEmpty must propagate through recursive depth-2 + // unnesting: an outer NULL or empty produces one NULL output row at + // each level. Adapted from `test_build_batch_list_arr_recursive`. + #[test] + fn test_build_batch_preserve_and_expand_empty_recursive() -> Result<()> { + // col1 | col2 + // [[1,2,3],null,[4,5]] | ['a','b'] + // [[7,8,9,10], null, [11,12,13]] | ['c','d'] + // null | ['e'] + let list_arr1 = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + let list_arr1_ref = Arc::new(list_arr1) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_null(); + let col1_field = Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + ))), + true, + ); + let col1 = ListArray::new( + Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr1_ref, + nulls.finish(), + ); + + let list_arr2 = StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ]); + let offsets = OffsetBuffer::from_lengths([2, 2, 1]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_n_non_nulls(3); + let col2_field = Field::new( + "col2", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ); + let col2 = GenericListArray::<i32>::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(list_arr2), + nulls.finish(), + ); + let schema = Arc::new(Schema::new(vec![col1_field, col2_field])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new( + "col1_unnest_placeholder_depth_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new("col1_unnest_placeholder_depth_2", DataType::Int32, true), + Field::new("col2_unnest_placeholder_depth_1", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(col1) as ArrayRef, Arc::new(col2) as ArrayRef], + )?; + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 0, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + None, + )? + .unwrap(); + + // The third input row (col1 = null, col2 = ['e']) now produces a + // NULL row for the depth-1 col1 placeholder *and* the depth-2 one, + // instead of being dropped at depth 1 and again at depth 2 the way + // it would be under `Drop`. Inner NULLs inside [...null...] sub- + // lists are still padded with NULL as before. + assert_snapshot!(batches_to_string(&[ret]), + @r" + +---------------------------------+---------------------------------+---------------------------------+ + | col1_unnest_placeholder_depth_1 | col1_unnest_placeholder_depth_2 | col2_unnest_placeholder_depth_1 | + +---------------------------------+---------------------------------+---------------------------------+ + | [1, 2, 3] | 1 | a | + | | 2 | b | + | [4, 5] | 3 | | + | [1, 2, 3] | | a | + | | | b | + | [4, 5] | | | + | [1, 2, 3] | 4 | a | + | | 5 | b | + | [4, 5] | | | + | [7, 8, 9, 10] | 7 | c | + | | 8 | d | + | [11, 12, 13] | 9 | | + | | 10 | | + | [7, 8, 9, 10] | | c | + | | | d | + | [11, 12, 13] | | | + | [7, 8, 9, 10] | 11 | c | + | | 12 | d | + | [11, 12, 13] | 13 | | + | | | e | + +---------------------------------+---------------------------------+---------------------------------+ + "); + Ok(()) + } + + #[test] + fn test_unnest_list_array() -> Result<()> { + // [A, B, C], [], NULL, [D], NULL, [NULL, F] + let list_array = make_generic_array::<i32>(); + verify_unnest_list_array( + &list_array, + vec![3, 2, 1, 2, 0, 3], + vec![ + Some("A"), + Some("B"), + Some("C"), + None, + None, + None, + Some("D"), + None, + None, + Some("F"), + None, + ], + )?; + + // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] + let list_array = make_fixed_list(); + verify_unnest_list_array( + &list_array, + vec![3, 1, 2, 0, 2, 3], + vec![ + Some("A"), + Some("B"), + None, + None, + Some("C"), + Some("D"), + None, + Some("F"), + None, + None, + None, + ], + )?; + + Ok(()) + } + + fn verify_longest_length( + list_arrays: &[ArrayRef], + null_handling: NullHandling, + expected: Vec<i64>, + ) -> Result<()> { + let options = UnnestOptions { + null_handling, + recursions: vec![], + }; + let longest_length = find_longest_length(list_arrays, &options)?; + let expected_array = Int64Array::from(expected); + assert_eq!( + longest_length + .as_any() + .downcast_ref::<Int64Array>() + .unwrap(), + &expected_array + ); + Ok(()) + } + + #[test] + fn test_longest_list_length() -> Result<()> { + // Test with single ListArray + // [A, B, C], [], NULL, [D], NULL, [NULL, F] + let list_array = Arc::new(make_generic_array::<i32>()) as ArrayRef; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![3, 0, 1, 1, 1, 2], + )?; + // PreserveAndExpandEmpty also treats empty lists as a NULL row. + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 1, 1, 1, 2], + )?; + + // Test with single LargeListArray + // [A, B, C], [], NULL, [D], NULL, [NULL, F] + let list_array = Arc::new(make_generic_array::<i64>()) as ArrayRef; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![3, 0, 1, 1, 1, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 1, 1, 1, 2], + )?; + + // Test with single FixedSizeListArray + // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] + let list_array = Arc::new(make_fixed_list()) as ArrayRef; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![2, 0, 2, 0, 2, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![2, 1, 2, 1, 2, 2], + )?; + + // Test with multiple list arrays + // [A, B, C], [], NULL, [D], NULL, [NULL, F] + // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] + let list1 = Arc::new(make_generic_array::<i32>()) as ArrayRef; + let list2 = Arc::new(make_fixed_list()) as ArrayRef; + let list_arrays = vec![Arc::clone(&list1), Arc::clone(&list2)]; + verify_longest_length(&list_arrays, NullHandling::Drop, vec![3, 0, 2, 1, 2, 2])?; + verify_longest_length( + &list_arrays, + NullHandling::Preserve, + vec![3, 1, 2, 1, 2, 2], + )?; + verify_longest_length( + &list_arrays, + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 2, 1, 2, 2], + )?; + + Ok(()) + } + + #[test] + fn test_create_take_indices() -> Result<()> { + let length_array = Int64Array::from(vec![2, 3, 1]); + let take_indices = create_take_indices(&length_array, 6); + let expected = Int64Array::from(vec![0, 0, 1, 1, 1, 2]); + assert_eq!(take_indices, expected); + Ok(()) + } + + /// Build a single-column `List<Int32>` batch where row `i` holds `lens[i]` elements, + /// numbered consecutively from 0 across the whole batch. A `None` length is a NULL + /// list. + fn list_batch(lens: &[Option<usize>]) -> RecordBatch { + let mut next = 0i32; + let rows: Vec<Option<Vec<Option<i32>>>> = lens + .iter() + .map(|len| { + len.map(|len| { + (0..len) + .map(|_| { + next += 1; + Some(next - 1) + }) + .collect() + }) + }) + .collect(); + let list = ListArray::from_iter_primitive::<Int32Type, _, _>(rows); + let schema = Arc::new(Schema::new(vec![Field::new( + "l", + list.data_type().clone(), + true, + )])); + RecordBatch::try_new(schema, vec![Arc::new(list)]).unwrap() + } + + /// Run a depth-1 unnest of column "l" over `input`, with the given + /// `datafusion.execution.batch_size`, and return the output batches. + async fn unnest_with_batch_size( + input: Vec<RecordBatch>, + batch_size: usize, + options: UnnestOptions, + ) -> Result<Vec<RecordBatch>> { + unnest_at_depth(input, batch_size, options, 1).await + } + + /// Unnest column "l" of `input` to `depth`, with the given + /// `datafusion.execution.batch_size`, and return the output batches. + async fn unnest_at_depth( + input: Vec<RecordBatch>, + batch_size: usize, + options: UnnestOptions, + depth: usize, + ) -> Result<Vec<RecordBatch>> { + let input_schema = input[0].schema(); + let output_schema = + Arc::new(Schema::new(vec![Field::new("l", DataType::Int32, true)])); + let source = + crate::test::TestMemoryExec::try_new_exec(&[input], input_schema, None)?; + let unnest = UnnestExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 0, + depth, + }], + vec![], + output_schema, + options, + )?; + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + datafusion_execution::config::SessionConfig::new() + .with_batch_size(batch_size), + ), + ); + crate::common::collect(unnest.execute(0, task_ctx)?).await + } + + /// The values an unnest produces, flattened across all output batches. + fn output_values(batches: &[RecordBatch]) -> Vec<Option<i32>> { + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::<Int32Type>() + .iter() + .collect::<Vec<_>>() + }) + .collect() + } + + /// Output batch sizes are fully determined by the input lengths and `batch_size`, so + /// assert the exact shapes rather than just the `<= batch_size` bound. Each case pins a + /// distinct path through `next_chunk_rows`. + #[tokio::test] + async fn test_unnest_stream_output_batch_shapes() -> Result<()> { + struct Case { + /// One inner slice per input batch, each of that batch's per-row list lengths. + lens_per_batch: &'static [&'static [Option<usize>]], + batch_size: usize, + expected_sizes: &'static [usize], + } + let cases: &[Case] = &[ + // Chunks pack several input rows. This is the case that distinguishes chunking + // the input from building everything and slicing: slicing a single 30-row build + // would give [8, 8, 8, 6]. + Case { + lens_per_batch: &[&[Some(3); 10]], + batch_size: 8, + expected_sizes: &[6, 6, 6, 6, 6], + }, + // Output smaller than batch_size comes back as one batch. + Case { + lens_per_batch: &[&[Some(3), Some(2)]], + batch_size: 1024, + expected_sizes: &[5], + }, + // One row expanding past batch_size cannot be chunked on the input side, so the + // oversized build is sliced on the way out instead. + Case { + lens_per_batch: &[&[Some(25)]], + batch_size: 10, + expected_sizes: &[10, 10, 5], + }, + // Chunk boundaries are per input batch, so each batch contributes a short tail. + Case { + lens_per_batch: &[&[Some(5), Some(5)], &[Some(1)], &[Some(7), Some(2)]], + batch_size: 4, + expected_sizes: &[4, 1, 4, 1, 1, 4, 3, 2], + }, + ]; + + for case in cases { + let input: Vec<RecordBatch> = case + .lens_per_batch + .iter() + .map(|lens| list_batch(lens)) + .collect(); + let batches = + unnest_with_batch_size(input, case.batch_size, UnnestOptions::default()) + .await?; + + let sizes: Vec<usize> = batches.iter().map(|b| b.num_rows()).collect(); + assert_eq!( + sizes, case.expected_sizes, + "lens={:?} batch_size={}", + case.lens_per_batch, case.batch_size + ); + + // `list_batch` numbers each batch's elements from 0, so the expected values are + // one run per input batch. Splitting must not perturb values or their order. + let expected_values: Vec<Option<i32>> = case + .lens_per_batch + .iter() + .flat_map(|lens| { + (0..lens.iter().flatten().sum::<usize>() as i32).map(Some) + }) + .collect(); + assert_eq!( + output_values(&batches), + expected_values, + "lens={:?} batch_size={}", + case.lens_per_batch, + case.batch_size + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_unnest_stream_chunking_preserves_null_handling() -> Result<()> { + // NULL and empty lists each contribute one NULL output row under + // PreserveAndExpandEmpty, and the per-row output counts that drive chunking must + // agree with that or chunk boundaries would drift out of step with the unnesting. + let lens = &[Some(3), Some(0), None, Some(2), None, Some(0)]; + let options = + UnnestOptions::new().with_null_handling(NullHandling::PreserveAndExpandEmpty); + + let chunked = + unnest_with_batch_size(vec![list_batch(lens)], 2, options.clone()).await?; + let whole = unnest_with_batch_size(vec![list_batch(lens)], 1024, options).await?; + + assert!(chunked.iter().all(|b| b.num_rows() <= 2)); + // 3 + 1 + 1 + 2 + 1 + 1 + assert_eq!(chunked.iter().map(|b| b.num_rows()).sum::<usize>(), 9); + assert_eq!(output_values(&chunked), output_values(&whole)); + Ok(()) + } + + #[tokio::test] + async fn test_unnest_stream_drop_null_handling() -> Result<()> { + // Under Drop, NULL and empty lists produce nothing. Chunks made up entirely of + // such rows yield no batch at all, and must not stall the stream or leak an + // empty batch into the output. + let lens = &[None, Some(0), None, Some(4), Some(0), None]; + let options = UnnestOptions::new().with_null_handling(NullHandling::Drop); + + let batches = unnest_with_batch_size(vec![list_batch(lens)], 2, options).await?; + + assert!(batches.iter().all(|b| b.num_rows() > 0)); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 4); + Ok(()) + } + + #[tokio::test] + async fn test_unnest_stream_recursive_respects_batch_size() -> Result<()> { + // Recursive unnest cannot have its expansion predicted from the input, so it falls + // back to unnesting a whole input batch and slicing the output. The batch_size + // guarantee has to hold on that path too. + let inner = Field::new_list_field(DataType::Int32, true); + let outer = + Field::new_list_field(DataType::new_list(DataType::Int32, true), true); + let values = Int32Array::from((0..24).collect::<Vec<_>>()); + // 12 inner lists of 2 elements each... + let inner_list = ListArray::new( + Arc::new(inner), + OffsetBuffer::new((0..=12).map(|i| i * 2).collect::<Vec<i32>>().into()), + Arc::new(values), + None, + ); + // ...grouped 3 to a row, so 4 input rows expand to 24 output rows at depth 2. + let outer_list = ListArray::new( + Arc::new(outer), + OffsetBuffer::new((0..=4).map(|i| i * 3).collect::<Vec<i32>>().into()), + Arc::new(inner_list), + None, + ); + let input_schema = Arc::new(Schema::new(vec![Field::new( + "l", + outer_list.data_type().clone(), + true, + )])); + let input = RecordBatch::try_new(input_schema, vec![Arc::new(outer_list)])?; + + let batches = + unnest_at_depth(vec![input], 7, UnnestOptions::default(), 2).await?; + + let sizes: Vec<usize> = batches.iter().map(|b| b.num_rows()).collect(); + assert_eq!(sizes, vec![7, 7, 7, 3]); + assert_eq!( + output_values(&batches), + (0..24).map(Some).collect::<Vec<_>>() + ); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/visitor.rs b/native/vendor/datafusion-physical-plan/src/visitor.rs new file mode 100644 index 0000000..892e603 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/visitor.rs
@@ -0,0 +1,94 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::ExecutionPlan; + +/// Visit all children of this plan, according to the order defined on `ExecutionPlanVisitor`. +// Note that this would be really nice if it were a method on +// ExecutionPlan, but it can not be because it takes a generic +// parameter and `ExecutionPlan` is a trait +pub fn accept<V: ExecutionPlanVisitor>( + plan: &dyn ExecutionPlan, + visitor: &mut V, +) -> Result<(), V::Error> { + visitor.pre_visit(plan)?; + for child in plan.children() { + visit_execution_plan(child.as_ref(), visitor)?; + } + visitor.post_visit(plan)?; + Ok(()) +} + +/// Trait that implements the [Visitor +/// pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for a +/// depth first walk of `ExecutionPlan` nodes. `pre_visit` is called +/// before any children are visited, and then `post_visit` is called +/// after all children have been visited. +/// +/// To use, define a struct that implements this trait and then invoke +/// ['accept']. +/// +/// For example, for an execution plan that looks like: +/// +/// ```text +/// ProjectionExec: id +/// FilterExec: state = CO +/// DataSourceExec: +/// ``` +/// +/// The sequence of visit operations would be: +/// ```text +/// visitor.pre_visit(ProjectionExec) +/// visitor.pre_visit(FilterExec) +/// visitor.pre_visit(DataSourceExec) +/// visitor.post_visit(DataSourceExec) +/// visitor.post_visit(FilterExec) +/// visitor.post_visit(ProjectionExec) +/// ``` +pub trait ExecutionPlanVisitor { + /// The type of error returned by this visitor + type Error; + + /// Invoked on an `ExecutionPlan` plan before any of its child + /// inputs have been visited. If Ok(true) is returned, the + /// recursion continues. If Err(..) or Ok(false) are returned, the + /// recursion stops immediately and the error, if any, is returned + /// to `accept` + fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result<bool, Self::Error>; + + /// Invoked on an `ExecutionPlan` plan *after* all of its child + /// inputs have been visited. The return value is handled the same + /// as the return value of `pre_visit`. The provided default + /// implementation returns `Ok(true)`. + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result<bool, Self::Error> { + Ok(true) + } +} + +/// Recursively calls `pre_visit` and `post_visit` for this node and +/// all of its children, as described on [`ExecutionPlanVisitor`] +pub fn visit_execution_plan<V: ExecutionPlanVisitor>( + plan: &dyn ExecutionPlan, + visitor: &mut V, +) -> Result<(), V::Error> { + visitor.pre_visit(plan)?; + for child in plan.children() { + visit_execution_plan(child.as_ref(), visitor)?; + } + visitor.post_visit(plan)?; + Ok(()) +}
diff --git a/native/vendor/datafusion-physical-plan/src/windows/bounded_window_agg_exec.rs b/native/vendor/datafusion-physical-plan/src/windows/bounded_window_agg_exec.rs new file mode 100644 index 0000000..d4c9800 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/windows/bounded_window_agg_exec.rs
@@ -0,0 +1,3030 @@ +// 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. + +//! Stream and channel implementations for window function expressions. +//! The executor given here uses bounded memory (does not maintain all +//! the input data seen so far), which makes it appropriate when processing +//! infinite inputs. + +use std::cmp::{Ordering, min}; +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use super::utils::create_schema; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::EmptyRecordBatchStream; +use crate::windows::{ + calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, + window_equivalence_properties, +}; +use crate::{ + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, + InputOrderMode, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, WindowExpr, validate_child_count, +}; + +use arrow::compute::take_record_batch; +use arrow::{ + array::{Array, ArrayRef, RecordBatchOptions, UInt32Array, UInt32Builder}, + compute::{concat, concat_batches, sort_to_indices, take_arrays}, + datatypes::SchemaRef, + record_batch::RecordBatch, +}; +use datafusion_common::hash_utils::create_hashes; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::{ + evaluate_partition_ranges, get_at_indices, get_row_at_idx, +}; +use datafusion_common::{ + HashMap, Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err, +}; +use datafusion_execution::TaskContext; +use datafusion_expr::ColumnarValue; +use datafusion_expr::window_state::{PartitionBatchState, WindowAggState}; +use datafusion_physical_expr::window::{ + PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowEvalContext, + WindowState, +}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::sort_expr::{ + OrderingRequirements, PhysicalSortExpr, +}; + +use crate::execution_plan::CardinalityEffect; +use datafusion_common::hash_utils::RandomState; +use futures::stream::Stream; +use futures::{StreamExt, ready}; +use hashbrown::hash_table::HashTable; +use indexmap::IndexMap; +use log::debug; + +/// Callback receiver for per-partition window state. +/// +/// `state` is the result of [`Accumulator::state`], which is a `&mut self` +/// call whose trait doc states "this function should not be called twice." +/// Several built-in aggregates (`median`, `percentile_cont`, `string_agg`, +/// `min_max_bytes`/`min_max_struct`) `std::mem::take` their internal +/// buffers to build that state — so `state` is a destructive read, not a +/// snapshot. The exec fires this at most once per group; a callee that +/// needs the value beyond the callback must retain it (e.g. clone into +/// owned storage). +/// +/// [`Accumulator::state`]: datafusion_expr::Accumulator::state +pub trait WindowStateObserver: Send + Sync { + /// Invoked once per (output-partition-index, window-expression, + /// PARTITION BY tuple) as each PARTITION BY group closes, for every + /// aggregate window expression on the exec. Non-aggregate window + /// functions (e.g. `row_number`, `rank`, `lead`/`lag`) do not fire this + /// callback. + /// + /// # Arguments + /// + /// * `partition_idx` - Output partition index of the [`BoundedWindowAggExec`] + /// stream firing this callback. + /// * `window_expr` - The window expression whose state just closed. + /// * `partition_key` - The PARTITION BY tuple that just closed. + /// * `state` - [`Accumulator::state`] for the closed group of + /// `window_expr`. See the trait-level doc for the destructive-read + /// contract. + /// + /// [`Accumulator::state`]: datafusion_expr::Accumulator::state + fn finalize_window_aggregate( + &self, + partition_idx: usize, + window_expr: &Arc<dyn WindowExpr>, + partition_key: &PartitionKey, + state: Vec<ScalarValue>, + ) -> Result<()>; +} + +/// Window execution plan +#[derive(Clone)] +pub struct BoundedWindowAggExec { + /// Input plan + input: Arc<dyn ExecutionPlan>, + /// Window function expression + window_expr: Vec<Arc<dyn WindowExpr>>, + /// Schema after the window is run + schema: SchemaRef, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Describes how the input is ordered relative to the partition keys + pub input_order_mode: InputOrderMode, + /// Partition by indices that define ordering + // For example, if input ordering is ORDER BY a, b and window expression + // contains PARTITION BY b, a; `ordered_partition_by_indices` would be 1, 0. + // Similarly, if window expression contains PARTITION BY a, b; then + // `ordered_partition_by_indices` would be 0, 1. + // See `get_ordered_partition_by_indices` for more details. + ordered_partition_by_indices: Vec<usize>, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// If `can_rerepartition` is false, partition_keys is always empty. + can_repartition: bool, + /// Invoked at partition-close to publish finalized per-partition window + /// state. Storage and multi-group handling are the caller's; the exec is + /// a pure event source. + state_observer: Option<Arc<dyn WindowStateObserver>>, +} + +impl std::fmt::Debug for BoundedWindowAggExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BoundedWindowAggExec") + .field("input", &self.input) + .field("window_expr", &self.window_expr) + .field("schema", &self.schema) + .field("metrics", &self.metrics) + .field("input_order_mode", &self.input_order_mode) + .field( + "ordered_partition_by_indices", + &self.ordered_partition_by_indices, + ) + .field("cache", &self.cache) + .field("can_repartition", &self.can_repartition) + .field( + "state_observer", + &self.state_observer.as_ref().map(|_| "..."), + ) + .finish() + } +} + +impl BoundedWindowAggExec { + /// Create a new execution plan for window aggregates + pub fn try_new( + window_expr: Vec<Arc<dyn WindowExpr>>, + input: Arc<dyn ExecutionPlan>, + input_order_mode: InputOrderMode, + can_repartition: bool, + ) -> Result<Self> { + let schema = create_schema(&input.schema(), &window_expr)?; + let schema = Arc::new(schema); + let partition_by_exprs = window_expr[0].partition_by(); + let ordered_partition_by_indices = match &input_order_mode { + InputOrderMode::Sorted => { + let indices = get_ordered_partition_by_indices( + window_expr[0].partition_by(), + &input, + )?; + if indices.len() == partition_by_exprs.len() { + indices + } else { + (0..partition_by_exprs.len()).collect::<Vec<_>>() + } + } + InputOrderMode::PartiallySorted(ordered_indices) => ordered_indices.clone(), + InputOrderMode::Linear => { + vec![] + } + }; + let cache = Self::compute_properties(&input, &schema, &window_expr)?; + Ok(Self { + input, + window_expr, + schema, + metrics: ExecutionPlanMetricsSet::new(), + input_order_mode, + ordered_partition_by_indices, + cache: Arc::new(cache), + can_repartition, + state_observer: None, + }) + } + + /// Install (or clear) a [`WindowStateObserver`] that receives each + /// PARTITION BY group's finalized window state at partition close. + /// + /// Errors when `observer` is `Some` and any window expression on this + /// exec has a non-ever-expanding frame (i.e. its start bound is not + /// `UNBOUNDED PRECEDING`). Those frames use `SlidingAggregateWindowExpr` + /// under the hood, whose accumulator calls `retract_batch` — at + /// partition close the accumulator holds only the last frame's rows, + /// not the partition aggregate, so the observed state would silently + /// misrepresent the group. + pub fn with_state_observer( + mut self, + observer: Option<Arc<dyn WindowStateObserver>>, + ) -> Result<Self> { + if observer.is_some() { + for expr in &self.window_expr { + if !expr.get_window_frame().is_ever_expanding() { + return exec_err!( + "cannot install WindowStateObserver on BoundedWindowAggExec \ + with a sliding aggregate window frame (start != \ + UNBOUNDED PRECEDING) for `{}`; sliding accumulator state \ + is frame-only, not the partition aggregate", + expr.name() + ); + } + } + } + self.state_observer = observer; + Ok(self) + } + + /// The currently-installed [`WindowStateObserver`], if any. Optimizer + /// rules that rebuild this exec via + /// [`crate::windows::get_best_fitting_window`] or a direct `try_new` + /// call must read this and reinstall it on the new exec, otherwise a + /// caller-installed observer is silently dropped by the rewrite. + pub fn state_observer(&self) -> Option<&Arc<dyn WindowStateObserver>> { + self.state_observer.as_ref() + } + + /// Window expressions + pub fn window_expr(&self) -> &[Arc<dyn WindowExpr>] { + &self.window_expr + } + + /// Input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Return the output sort order of partition keys: For example + /// OVER(PARTITION BY a, ORDER BY b) -> would give sorting of the column a + // We are sure that partition by columns are always at the beginning of sort_keys + // Hence returned `PhysicalSortExpr` corresponding to `PARTITION BY` columns can be used safely + // to calculate partition separation points + pub fn partition_by_sort_keys(&self) -> Result<Vec<PhysicalSortExpr>> { + let partition_by = self.window_expr()[0].partition_by(); + get_partition_by_sort_exprs( + &self.input, + partition_by, + &self.ordered_partition_by_indices, + ) + } + + /// Initializes the appropriate [`PartitionSearcher`] implementation from + /// the state. + fn get_search_algo(&self) -> Result<Box<dyn PartitionSearcher>> { + let partition_by_sort_keys = self.partition_by_sort_keys()?; + let ordered_partition_by_indices = self.ordered_partition_by_indices.clone(); + let input_schema = self.input().schema(); + Ok(match &self.input_order_mode { + InputOrderMode::Sorted => { + // In Sorted mode, all partition by columns should be ordered. + if self.window_expr()[0].partition_by().len() + != ordered_partition_by_indices.len() + { + return exec_err!( + "All partition by columns should have an ordering in Sorted mode." + ); + } + Box::new(SortedSearch { + partition_by_sort_keys, + ordered_partition_by_indices, + input_schema, + }) + } + InputOrderMode::Linear | InputOrderMode::PartiallySorted(_) => Box::new( + LinearSearch::new(ordered_partition_by_indices, input_schema), + ), + }) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + input: &Arc<dyn ExecutionPlan>, + schema: &SchemaRef, + window_exprs: &[Arc<dyn WindowExpr>], + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let eq_properties = window_equivalence_properties(schema, input, window_exprs)?; + + // As we can have repartitioning using the partition keys, this can + // be either one or more than one, depending on the presence of + // repartitioning. + let output_partitioning = input.output_partitioning().clone(); + + // Construct properties cache + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + // TODO: Emission type and boundedness information can be enhanced here + input.pipeline_behavior(), + input.boundedness(), + )) + } + + pub fn partition_keys(&self) -> Vec<Arc<dyn PhysicalExpr>> { + if !self.can_repartition { + vec![] + } else { + let all_partition_keys = self + .window_expr() + .iter() + .map(|expr| expr.partition_by().to_vec()) + .collect::<Vec<_>>(); + + all_partition_keys + .into_iter() + .min_by_key(|s| s.len()) + .unwrap_or_else(Vec::new) + } + } + + fn statistics_helper(&self, statistics: Statistics) -> Result<Statistics> { + let win_cols = self.window_expr.len(); + let input_cols = self.input.schema().fields().len(); + // TODO stats: some windowing function will maintain invariants such as min, max... + let mut column_statistics = Vec::with_capacity(win_cols + input_cols); + // copy stats of the input to the beginning of the schema. + column_statistics.extend(statistics.column_statistics); + for _ in 0..win_cols { + column_statistics.push(ColumnStatistics::new_unknown()) + } + Ok(Statistics { + num_rows: statistics.num_rows, + column_statistics, + total_byte_size: Precision::Absent, + }) + } +} + +impl DisplayAs for BoundedWindowAggExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "BoundedWindowAggExec: ")?; + let g: Vec<String> = self + .window_expr + .iter() + .map(|e| { + let field = match e.field() { + Ok(f) => f.to_string(), + Err(e) => format!("{e:?}"), + }; + format!( + "{}: {}, frame: {}", + e.name().to_owned(), + field, + e.get_window_frame() + ) + }) + .collect(); + let mode = &self.input_order_mode; + write!(f, "wdw=[{}], mode=[{:?}]", g.join(", "), mode)?; + } + DisplayFormatType::TreeRender => { + let g: Vec<String> = self + .window_expr + .iter() + .map(|e| e.name().to_owned().to_string()) + .collect(); + writeln!(f, "select_list={}", g.join(", "))?; + + let mode = &self.input_order_mode; + writeln!(f, "mode={mode:?}")?; + } + } + Ok(()) + } +} + +impl ExecutionPlan for BoundedWindowAggExec { + fn name(&self) -> &'static str { + "BoundedWindowAggExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let expressions = self.window_expr.iter().flat_map(|window_expr| { + let expressions = window_expr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.partition_by_exprs) + .chain(expressions.order_by_exprs) + }); + crate::apply_expression_roots(expressions, f) + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + let partition_bys = self.window_expr()[0].partition_by(); + let order_keys = self.window_expr()[0].order_by(); + let partition_bys = self + .ordered_partition_by_indices + .iter() + .map(|idx| &partition_bys[*idx]); + vec![calc_requirements(partition_bys, order_keys)] + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.partition_keys().is_empty() { + debug!("No partition defined for BoundedWindowAggExec!!!"); + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) + } else { + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) + } + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new = BoundedWindowAggExec::try_new( + self.window_expr.clone(), + Arc::clone(&children[0]), + self.input_order_mode.clone(), + self.can_repartition, + )? + .with_state_observer(self.state_observer.clone())?; + Ok(Arc::new(new)) + } + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, context)?; + let search_mode = self.get_search_algo()?; + let stream = Box::pin(BoundedWindowAggStream::new( + Arc::clone(&self.schema), + self.window_expr.clone(), + input, + BaselineMetrics::new(&self.metrics, partition), + search_mode, + partition, + self.state_observer.clone(), + )?); + Ok(stream) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let input_stat = input_stats[0].as_ref().clone(); + Ok(Arc::new(self.statistics_helper(input_stat)?)) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use super::proto::encode_physical_window_expr; + use datafusion_proto_common::protobuf_common::EmptyMessage; + use datafusion_proto_models::protobuf; + use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode; + + // Exhaustive destructure: adding a field to `BoundedWindowAggExec` + // without deciding how it is serialized is a compile error, not a + // silent round-trip gap. + let Self { + input, + window_expr, + // Derived at construction by `create_schema` from the input schema + // and the window expressions. + schema: _, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + input_order_mode, + // Derived at construction from `input_order_mode` and the window + // expressions' PARTITION BY. + ordered_partition_by_indices: _, + // Derived at construction by `Self::compute_properties`. + cache: _, + // No wire field of its own; it is folded into `partition_keys` + // below, since `partition_keys()` returns an empty vec when this is + // false and the decoder recovers it as `!partition_keys.is_empty()`. + can_repartition: _, + // Runtime callback installed after planning; not part of the wire + // format. Any decoder that needs it must reinstall via + // `with_state_observer`. + state_observer: _, + } = self; + + let input = ctx.encode_child(input)?; + let window_expr = window_expr + .iter() + .map(|expr| encode_physical_window_expr(expr, ctx)) + .collect::<Result<Vec<_>>>()?; + let partition_keys = self + .partition_keys() + .iter() + .map(|expr| ctx.encode_expr(expr)) + .collect::<Result<Vec<_>>>()?; + // A `Some(input_order_mode)` is what tells the shared `Window` decode + // arm to rebuild a `BoundedWindowAggExec` rather than a `WindowAggExec`. + let input_order_mode = match input_order_mode { + InputOrderMode::Linear => ProtoInputOrderMode::Linear(EmptyMessage {}), + InputOrderMode::PartiallySorted(columns) => { + ProtoInputOrderMode::PartiallySorted( + protobuf::PartiallySortedInputOrderMode { + columns: columns.iter().map(|column| *column as u64).collect(), + }, + ) + } + InputOrderMode::Sorted => ProtoInputOrderMode::Sorted(EmptyMessage {}), + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new( + protobuf::WindowAggExecNode { + input: Some(Box::new(input)), + window_expr, + partition_keys, + input_order_mode: Some(input_order_mode), + }, + )), + ), + })) + } +} + +/// Trait that specifies how we search for (or calculate) partitions. It has two +/// implementations: [`SortedSearch`] and [`LinearSearch`]. +trait PartitionSearcher: Send { + /// This method constructs output columns using the result of each window expression + /// (each entry in the output vector comes from a window expression). + /// Executor when producing output concatenates `input_buffer` (corresponding section), and + /// result of this function to generate output `RecordBatch`. `input_buffer` is used to determine + /// which sections of the window expression results should be used to generate output. + /// `partition_buffers` contains corresponding section of the `RecordBatch` for each partition. + /// `window_agg_states` stores per partition state for each window expression. + /// None case means that no result is generated + /// `Some(Vec<ArrayRef>)` is the result of each window expression. + fn calculate_out_columns( + &mut self, + input_buffer: &RecordBatch, + window_agg_states: &[PartitionWindowAggStates], + partition_buffers: &mut PartitionBatches, + window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Option<Vec<ArrayRef>>>; + + /// Determine whether `[InputOrderMode]` is `[InputOrderMode::Linear]` or not. + fn is_mode_linear(&self) -> bool { + false + } + + // Constructs corresponding batches for each partition for the record_batch. + fn evaluate_partition_batches( + &mut self, + record_batch: &RecordBatch, + window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Vec<(PartitionKey, RecordBatch)>>; + + /// Prunes the state. + fn prune(&mut self, _n_out: usize) {} + + /// Marks the partition as done if we are sure that corresponding partition + /// cannot receive any more values. + fn mark_partition_end(&self, partition_buffers: &mut PartitionBatches); + + /// Updates `input_buffer` and `partition_buffers` with the new `record_batch`. + fn update_partition_batch( + &mut self, + input_buffer: &mut RecordBatch, + record_batch: RecordBatch, + window_expr: &[Arc<dyn WindowExpr>], + partition_buffers: &mut PartitionBatches, + ) -> Result<()> { + if record_batch.num_rows() == 0 { + return Ok(()); + } + let partition_batches = + self.evaluate_partition_batches(&record_batch, window_expr)?; + for (partition_row, partition_batch) in partition_batches { + if let Some(partition_batch_state) = partition_buffers.get_mut(&partition_row) + { + partition_batch_state.extend(&partition_batch)? + } else { + let options = RecordBatchOptions::new() + .with_row_count(Some(partition_batch.num_rows())); + // Use input_schema for the buffer schema, not `record_batch.schema()` + // as it may not have the "correct" schema in terms of output + // nullability constraints. For details, see the following issue: + // https://github.com/apache/datafusion/issues/9320 + let partition_batch = RecordBatch::try_new_with_options( + Arc::clone(self.input_schema()), + partition_batch.columns().to_vec(), + &options, + )?; + let partition_batch_state = + PartitionBatchState::new_with_batch(partition_batch); + partition_buffers.insert(partition_row, partition_batch_state); + } + } + + self.mark_partition_end(partition_buffers); + + *input_buffer = if input_buffer.num_rows() == 0 { + record_batch + } else { + concat_batches(self.input_schema(), [input_buffer, &record_batch])? + }; + + Ok(()) + } + + fn input_schema(&self) -> &SchemaRef; +} + +/// This object encapsulates the algorithm state for a simple linear scan +/// algorithm for computing partitions. +pub struct LinearSearch { + /// Keeps the hash of input buffer calculated from PARTITION BY columns. + /// Its length is equal to the `input_buffer` length. + input_buffer_hashes: VecDeque<u64>, + /// Used during hash value calculation. + random_state: RandomState, + /// Input ordering and partition by key ordering need not be the same, so + /// this vector stores the mapping between them. For instance, if the input + /// is ordered by a, b and the window expression contains a PARTITION BY b, a + /// clause, this attribute stores [1, 0]. + ordered_partition_by_indices: Vec<usize>, + /// We use this [`HashTable`] to calculate unique partitions for each new + /// RecordBatch. First entry in the tuple is the hash value, the second + /// entry is the unique ID for each partition (increments from 0 to n). + row_map_batch: HashTable<(u64, usize)>, + /// We use this [`HashTable`] to calculate the output columns that we can + /// produce at each cycle. First entry in the tuple is the hash value, the + /// second entry is the unique ID for each partition (increments from 0 to n). + /// The third entry stores how many new outputs are calculated for the + /// corresponding partition. + row_map_out: HashTable<(u64, usize, usize)>, + input_schema: SchemaRef, +} + +impl PartitionSearcher for LinearSearch { + /// This method constructs output columns using the result of each window expression. + // Assume input buffer is | Partition Buffers would be (Where each partition and its data is separated) + // a, 2 | a, 2 + // b, 2 | a, 2 + // a, 2 | a, 2 + // b, 2 | + // a, 2 | b, 2 + // b, 2 | b, 2 + // b, 2 | b, 2 + // | b, 2 + // Also assume we happen to calculate 2 new values for a, and 3 for b (To be calculate missing values we may need to consider future values). + // Partition buffers effectively will be + // a, 2, 1 + // a, 2, 2 + // a, 2, (missing) + // + // b, 2, 1 + // b, 2, 2 + // b, 2, 3 + // b, 2, (missing) + // When partition buffers are mapped back to the original record batch. Result becomes + // a, 2, 1 + // b, 2, 1 + // a, 2, 2 + // b, 2, 2 + // a, 2, (missing) + // b, 2, 3 + // b, 2, (missing) + // This function calculates the column result of window expression(s) (First 4 entry of 3rd column in the above section.) + // 1 + // 1 + // 2 + // 2 + // Above section corresponds to calculated result which can be emitted without breaking input buffer ordering. + fn calculate_out_columns( + &mut self, + input_buffer: &RecordBatch, + window_agg_states: &[PartitionWindowAggStates], + partition_buffers: &mut PartitionBatches, + window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Option<Vec<ArrayRef>>> { + let partition_output_indices = self.calc_partition_output_indices( + input_buffer, + window_agg_states, + window_expr, + )?; + + let n_window_col = window_agg_states.len(); + let mut new_columns = vec![vec![]; n_window_col]; + // Size of all_indices can be at most input_buffer.num_rows(): + let mut all_indices = UInt32Builder::with_capacity(input_buffer.num_rows()); + for (row, indices) in partition_output_indices { + let length = indices.len(); + for (idx, window_agg_state) in window_agg_states.iter().enumerate() { + let partition = &window_agg_state[&row]; + let values = Arc::clone(&partition.state.out_col.slice(0, length)); + new_columns[idx].push(values); + } + let partition_batch_state = &mut partition_buffers[&row]; + // Store how many rows are generated for each partition + partition_batch_state.n_out_row = length; + // For each row keep corresponding index in the input record batch + all_indices.append_slice(&indices); + } + let all_indices = all_indices.finish(); + if all_indices.is_empty() { + // We couldn't generate any new value, return early: + return Ok(None); + } + + // Concatenate results for each column by converting `Vec<Vec<ArrayRef>>` + // to Vec<ArrayRef> where inner `Vec<ArrayRef>`s are converted to `ArrayRef`s. + let new_columns = new_columns + .iter() + .map(|items| { + concat(&items.iter().map(|e| e.as_ref()).collect::<Vec<_>>()) + .map_err(|e| arrow_datafusion_err!(e)) + }) + .collect::<Result<Vec<_>>>()?; + // We should emit columns according to row index ordering. + let sorted_indices = sort_to_indices(&all_indices, None, None)?; + // Construct new column according to row ordering. This fixes ordering + take_arrays(&new_columns, &sorted_indices, None) + .map(Some) + .map_err(|e| arrow_datafusion_err!(e)) + } + + fn evaluate_partition_batches( + &mut self, + record_batch: &RecordBatch, + window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Vec<(PartitionKey, RecordBatch)>> { + let partition_bys = + evaluate_partition_by_column_values(record_batch, window_expr)?; + // NOTE: In Linear or PartiallySorted modes, we are sure that + // `partition_bys` are not empty. + let (mut keys, permutation, bounds) = + self.compute_partition_permutation(&partition_bys, record_batch)?; + if keys.len() == 1 { + // The batch contains a single partition, so the gather below + // would be an identity permutation; use the batch as-is. + let key = keys.remove(0); + return Ok(vec![(key, record_batch.clone())]); + } + // Reorder the batch with a single `take` so that each partition's + // rows become contiguous, then hand each partition a zero-copy slice + // of the result. The slices share the gathered batch's buffers; + // `PartitionBatchState::extend` copies out of them the next time the + // partition receives rows. + let gathered = take_record_batch(record_batch, &UInt32Array::from(permutation))?; + Ok(keys + .into_iter() + .zip(bounds.windows(2)) + .map(|(key, bound)| (key, gathered.slice(bound[0], bound[1] - bound[0]))) + .collect()) + } + + fn prune(&mut self, n_out: usize) { + // Delete hashes for the rows that are outputted. + self.input_buffer_hashes.drain(0..n_out); + } + + fn mark_partition_end(&self, partition_buffers: &mut PartitionBatches) { + // We should be in the `PartiallySorted` case, otherwise we can not + // tell when we are at the end of a given partition. + if !self.ordered_partition_by_indices.is_empty() + && let Some((last_row, _)) = partition_buffers.last() + { + let last_sorted_cols = self + .ordered_partition_by_indices + .iter() + .map(|idx| last_row[*idx].clone()) + .collect::<Vec<_>>(); + for (row, partition_batch_state) in partition_buffers.iter_mut() { + let sorted_cols = self + .ordered_partition_by_indices + .iter() + .map(|idx| &row[*idx]); + // All the partitions other than `last_sorted_cols` are done. + // We are sure that we will no longer receive values for these + // partitions (arrival of a new value would violate ordering). + partition_batch_state.is_end = !sorted_cols.eq(&last_sorted_cols); + } + } + } + + fn is_mode_linear(&self) -> bool { + self.ordered_partition_by_indices.is_empty() + } + + fn input_schema(&self) -> &SchemaRef { + &self.input_schema + } +} + +impl LinearSearch { + /// Initialize a new [`LinearSearch`] partition searcher. + fn new(ordered_partition_by_indices: Vec<usize>, input_schema: SchemaRef) -> Self { + LinearSearch { + input_buffer_hashes: VecDeque::new(), + random_state: Default::default(), + ordered_partition_by_indices, + row_map_batch: HashTable::with_capacity(256), + row_map_out: HashTable::with_capacity(256), + input_schema, + } + } + + /// Splits the rows of `batch` by partition, according to the PARTITION BY + /// expression results in `columns`. Returns the distinct partition keys + /// in first-appearance order, a permutation of the row indices of + /// `batch` that groups each partition's rows together, and the + /// boundaries of each partition's run of rows within that permutation: + /// partition `p` occupies `permutation[bounds[p]..bounds[p + 1]]`, and + /// its indices are in ascending (stream) order. + fn compute_partition_permutation( + &mut self, + columns: &[ArrayRef], + batch: &RecordBatch, + ) -> Result<(Vec<PartitionKey>, Vec<u32>, Vec<usize>)> { + let num_rows = batch.num_rows(); + let mut batch_hashes = vec![0; num_rows]; + create_hashes(columns, &self.random_state, &mut batch_hashes)?; + self.input_buffer_hashes.extend(&batch_hashes); + // reset row_map for new calculation + self.row_map_batch.clear(); + let mut keys: Vec<PartitionKey> = vec![]; + // Partition id of each row, in row order: + let mut row_partition_ids = Vec::with_capacity(num_rows); + // Number of rows in each partition: + let mut counts: Vec<usize> = vec![]; + for (hash, row_idx) in batch_hashes.into_iter().zip(0u32..) { + let entry = self.row_map_batch.find_mut(hash, |(_, group_idx)| { + let row = get_row_at_idx(columns, row_idx as usize).unwrap(); + // Handle hash collisions with an equality check: + row == keys[*group_idx] + }); + let group_idx = if let Some((_, group_idx)) = entry { + *group_idx + } else { + let group_idx = keys.len(); + self.row_map_batch + .insert_unique(hash, (hash, group_idx), |(hash, _)| *hash); + keys.push(get_row_at_idx(columns, row_idx as usize)?); + counts.push(0); + group_idx + }; + row_partition_ids.push(group_idx); + counts[group_idx] += 1; + } + // A prefix sum over the counts gives each partition's run boundaries + // in the permutation. + let mut bounds = Vec::with_capacity(counts.len() + 1); + let mut total = 0; + bounds.push(0); + for count in counts { + total += count; + bounds.push(total); + } + // Scatter each row's index into its partition's run. Visiting rows + // in ascending order keeps each run in ascending row order. + let mut cursors: Vec<usize> = bounds[..bounds.len() - 1].to_vec(); + let mut permutation = vec![0u32; num_rows]; + for (row_idx, group_idx) in row_partition_ids.into_iter().enumerate() { + permutation[cursors[group_idx]] = row_idx as u32; + cursors[group_idx] += 1; + } + Ok((keys, permutation, bounds)) + } + + /// Calculates partition keys and result indices for each partition. + /// The return value is a vector of tuples where the first entry stores + /// the partition key (unique for each partition) and the second entry + /// stores indices of the rows for which the partition is constructed. + fn calc_partition_output_indices( + &mut self, + input_buffer: &RecordBatch, + window_agg_states: &[PartitionWindowAggStates], + window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Vec<(PartitionKey, Vec<u32>)>> { + let partition_by_columns = + evaluate_partition_by_column_values(input_buffer, window_expr)?; + // Reset the row_map state: + self.row_map_out.clear(); + let mut partition_indices: Vec<(PartitionKey, Vec<u32>)> = vec![]; + for (hash, row_idx) in self.input_buffer_hashes.iter().zip(0u32..) { + let entry = self.row_map_out.find_mut(*hash, |(_, group_idx, _)| { + let row = + get_row_at_idx(&partition_by_columns, row_idx as usize).unwrap(); + row == partition_indices[*group_idx].0 + }); + if let Some((_, group_idx, n_out)) = entry { + let (_, indices) = &mut partition_indices[*group_idx]; + if indices.len() >= *n_out { + break; + } + indices.push(row_idx); + } else { + let row = get_row_at_idx(&partition_by_columns, row_idx as usize)?; + let min_out = window_agg_states + .iter() + .map(|window_agg_state| { + window_agg_state + .get(&row) + .map(|partition| partition.state.out_col.len()) + .unwrap_or(0) + }) + .min() + .unwrap_or(0); + if min_out == 0 { + break; + } + self.row_map_out.insert_unique( + *hash, + (*hash, partition_indices.len(), min_out), + |(hash, _, _)| *hash, + ); + partition_indices.push((row, vec![row_idx])); + } + } + Ok(partition_indices) + } +} + +/// This object encapsulates the algorithm state for sorted searching +/// when computing partitions. +pub struct SortedSearch { + /// Stores partition by columns and their ordering information + partition_by_sort_keys: Vec<PhysicalSortExpr>, + /// Input ordering and partition by key ordering need not be the same, so + /// this vector stores the mapping between them. For instance, if the input + /// is ordered by a, b and the window expression contains a PARTITION BY b, a + /// clause, this attribute stores [1, 0]. + ordered_partition_by_indices: Vec<usize>, + input_schema: SchemaRef, +} + +impl PartitionSearcher for SortedSearch { + /// This method constructs new output columns using the result of each window expression. + fn calculate_out_columns( + &mut self, + _input_buffer: &RecordBatch, + window_agg_states: &[PartitionWindowAggStates], + partition_buffers: &mut PartitionBatches, + _window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Option<Vec<ArrayRef>>> { + let n_out = self.calculate_n_out_row(window_agg_states, partition_buffers); + if n_out == 0 { + Ok(None) + } else { + window_agg_states + .iter() + .map(|map| get_aggregate_result_out_column(map, n_out).map(Some)) + .collect() + } + } + + fn evaluate_partition_batches( + &mut self, + record_batch: &RecordBatch, + _window_expr: &[Arc<dyn WindowExpr>], + ) -> Result<Vec<(PartitionKey, RecordBatch)>> { + let num_rows = record_batch.num_rows(); + // Calculate result of partition by column expressions + let partition_columns = self + .partition_by_sort_keys + .iter() + .map(|elem| elem.evaluate_to_sort_column(record_batch)) + .collect::<Result<Vec<_>>>()?; + // Reorder `partition_columns` such that its ordering matches input ordering. + let partition_columns_ordered = + get_at_indices(&partition_columns, &self.ordered_partition_by_indices)?; + let partition_points = + evaluate_partition_ranges(num_rows, &partition_columns_ordered)?; + let partition_bys = partition_columns + .into_iter() + .map(|arr| arr.values) + .collect::<Vec<ArrayRef>>(); + + partition_points + .iter() + .map(|range| { + let row = get_row_at_idx(&partition_bys, range.start)?; + let len = range.end - range.start; + let slice = record_batch.slice(range.start, len); + Ok((row, slice)) + }) + .collect::<Result<Vec<_>>>() + } + + fn mark_partition_end(&self, partition_buffers: &mut PartitionBatches) { + // In Sorted case. We can mark all partitions besides last partition as ended. + // We are sure that those partitions will never receive any values. + // (Otherwise ordering invariant is violated.) + let n_partitions = partition_buffers.len(); + for (idx, (_, partition_batch_state)) in partition_buffers.iter_mut().enumerate() + { + partition_batch_state.is_end |= idx < n_partitions - 1; + } + } + + fn input_schema(&self) -> &SchemaRef { + &self.input_schema + } +} + +impl SortedSearch { + /// Calculates how many rows we can output. + fn calculate_n_out_row( + &mut self, + window_agg_states: &[PartitionWindowAggStates], + partition_buffers: &mut PartitionBatches, + ) -> usize { + // Different window aggregators may produce results at different rates. + // We produce the overall batch result only as fast as the slowest one. + let mut counts = vec![]; + let out_col_counts = window_agg_states.iter().map(|window_agg_state| { + // Store how many elements are generated for the current + // window expression: + let mut cur_window_expr_out_result_len = 0; + // We iterate over `window_agg_state`, which is an IndexMap. + // Iterations follow the insertion order, hence we preserve + // sorting when partition columns are sorted. + let mut per_partition_out_results = HashMap::new(); + for (row, WindowState { state, .. }) in window_agg_state.iter() { + cur_window_expr_out_result_len += state.out_col.len(); + let count = per_partition_out_results.entry(row).or_insert(0); + if *count < state.out_col.len() { + *count = state.out_col.len(); + } + // If we do not generate all results for the current + // partition, we do not generate results for next + // partition -- otherwise we will lose input ordering. + if state.n_row_result_missing > 0 { + break; + } + } + counts.push(per_partition_out_results); + cur_window_expr_out_result_len + }); + argmin(out_col_counts).map_or(0, |(min_idx, minima)| { + let mut slowest_partition = counts.swap_remove(min_idx); + for (partition_key, partition_batch) in partition_buffers.iter_mut() { + if let Some(count) = slowest_partition.remove(partition_key) { + partition_batch.n_out_row = count; + } + } + minima + }) + } +} + +/// Calculates partition by expression results for each window expression +/// on `record_batch`. +fn evaluate_partition_by_column_values( + record_batch: &RecordBatch, + window_expr: &[Arc<dyn WindowExpr>], +) -> Result<Vec<ArrayRef>> { + window_expr[0] + .partition_by() + .iter() + .map(|item| match item.evaluate(record_batch)? { + ColumnarValue::Array(array) => Ok(array), + ColumnarValue::Scalar(scalar) => { + scalar.to_array_of_size(record_batch.num_rows()) + } + }) + .collect() +} + +/// Stream for the bounded window aggregation plan. +pub struct BoundedWindowAggStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + /// The record batch executor receives as input (i.e. the columns needed + /// while calculating aggregation results). + input_buffer: RecordBatch, + /// Each partition's rows, accumulated across input batches. All window + /// expressions calculate their results against these shared rows without + /// copying. + partition_buffers: PartitionBatches, + /// An executor can run multiple window expressions if the PARTITION BY + /// and ORDER BY sections are same. We keep state of the each window + /// expression inside `window_agg_states`. + window_agg_states: Vec<PartitionWindowAggStates>, + finished: bool, + window_expr: Vec<Arc<dyn WindowExpr>>, + baseline_metrics: BaselineMetrics, + /// Search mode for partition columns. This determines the algorithm with + /// which we group each partition. + search_mode: Box<dyn PartitionSearcher>, + /// In `Linear` mode, a single-row batch containing the most recent input + /// row (whichever partition that row belongs to); `None` in other modes + /// and before the first non-empty batch arrives. Since in `Linear` mode + /// the input is sorted by the first ORDER BY column, no future input row + /// -- in any partition -- can precede this row in that column. Every + /// partition's evaluation consults this bound to decide whether pending + /// window frames can be finalized before the partition receives more + /// data (which in turn allows buffered state to be pruned). Note that + /// only the first ORDER BY column provides this guarantee. As a counter + /// example, consider `PARTITION BY b, ORDER BY a, c` when the input is + /// sorted by `[a, b, c]`: the mode will be `Linear`, but the last row of + /// the input is the "last" data in terms of `[a, b, c]`, not in terms of + /// the ordering requirement `[a, c]`. Hence, only column `a` can serve + /// as a guarantee of the "last" data across partitions. In the `Sorted` + /// and `PartiallySorted` modes, the leading ordering separates + /// partitions, so finished partitions are pruned eagerly instead and no + /// such bound is needed. + most_recent_row: Option<RecordBatch>, + /// Output partition index this stream serves; passed as the first + /// argument to [`WindowStateObserver::finalize_window_aggregate`]. + partition_idx: usize, + /// If set, invoked from [`Self::publish_finalized_states`] with the + /// finalized per-window-expression state for every partition key that is + /// about to be dropped. + state_observer: Option<Arc<dyn WindowStateObserver>>, +} + +impl BoundedWindowAggStream { + /// Fire `observer` once per (window expression, partition key) for every + /// group whose [`WindowAggState::is_end`] is true. Always mutates when + /// called: [`datafusion_expr::Accumulator::state`] requires `&mut`, which + /// propagates up here. The caller is responsible for deciding whether to + /// fire (i.e. checking whether an observer is installed). + /// + /// Exactly-once per group is enforced by [`WindowState::aggregate_state`], + /// which errors on second call; the `published` early-skip below avoids reaching the error. + fn publish_finalized_states( + &mut self, + observer: &dyn WindowStateObserver, + ) -> Result<()> { + let partition_idx = self.partition_idx; + for (expr_idx, per_expr) in self.window_agg_states.iter_mut().enumerate() { + let window_expr = &self.window_expr[expr_idx]; + for (key, ws) in per_expr.iter_mut() { + if ws.published || !ws.state.is_end { + continue; + } + if let Some(state) = ws.aggregate_state()? { + observer.finalize_window_aggregate( + partition_idx, + window_expr, + key, + state, + )?; + } + } + } + Ok(()) + } + + /// Prunes sections of the state that are no longer needed when calculating + /// results (as determined by window frame boundaries and number of results generated). + // For instance, if first `n` (not necessarily same with `n_out`) elements are no longer needed to + // calculate window expression result (outside the window frame boundary) we retract first `n` elements + // from the corresponding partition's batch in `self.partition_buffers`. + // For instance, if `n_out` number of rows are calculated, we can remove + // first `n_out` rows from `self.input_buffer`. + fn prune_state(&mut self, n_out: usize) -> Result<()> { + // Prune `self.window_agg_states`: + self.prune_out_columns(); + // Prune `self.partition_buffers`: + self.prune_partition_batches(); + // Prune `self.input_buffer`: + self.prune_input_batch(n_out)?; + // Prune internal state of search algorithm. + self.search_mode.prune(n_out); + Ok(()) + } +} + +impl Stream for BoundedWindowAggStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let poll = self.poll_next_inner(cx); + self.baseline_metrics.record_poll(poll) + } +} + +impl BoundedWindowAggStream { + /// Create a new BoundedWindowAggStream + fn new( + schema: SchemaRef, + window_expr: Vec<Arc<dyn WindowExpr>>, + input: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, + search_mode: Box<dyn PartitionSearcher>, + partition_idx: usize, + state_observer: Option<Arc<dyn WindowStateObserver>>, + ) -> Result<Self> { + let state = window_expr.iter().map(|_| IndexMap::default()).collect(); + let empty_batch = RecordBatch::new_empty(Arc::clone(&schema)); + Ok(Self { + schema, + input, + input_buffer: empty_batch, + partition_buffers: IndexMap::default(), + window_agg_states: state, + finished: false, + window_expr, + baseline_metrics, + search_mode, + most_recent_row: None, + partition_idx, + state_observer, + }) + } + + fn compute_aggregates(&mut self) -> Result<Option<RecordBatch>> { + // calculate window cols + let eval_ctx = WindowEvalContext::default() + .with_most_recent_row(self.most_recent_row.as_ref()); + for (cur_window_expr, state) in + self.window_expr.iter().zip(&mut self.window_agg_states) + { + cur_window_expr.evaluate_stateful( + &self.partition_buffers, + state, + &eval_ctx, + )?; + } + + // Fire before `calculate_out_columns`: on causal frames every row + // already streamed out, so at EOS that call returns `None` and the + // prune path is skipped — the final partition would otherwise be + // dropped unobserved. + if let Some(observer) = self.state_observer.clone() { + self.publish_finalized_states(observer.as_ref())?; + } + + let schema = Arc::clone(&self.schema); + let window_expr_out = self.search_mode.calculate_out_columns( + &self.input_buffer, + &self.window_agg_states, + &mut self.partition_buffers, + &self.window_expr, + )?; + if let Some(window_expr_out) = window_expr_out { + let n_out = window_expr_out[0].len(); + // right append new columns to corresponding section in the original input buffer. + let columns_to_show = self + .input_buffer + .columns() + .iter() + .map(|elem| elem.slice(0, n_out)) + .chain(window_expr_out) + .collect::<Vec<_>>(); + let n_generated = columns_to_show[0].len(); + self.prune_state(n_generated)?; + Ok(Some(RecordBatch::try_new(schema, columns_to_show)?)) + } else { + Ok(None) + } + } + + #[inline] + fn poll_next_inner( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + if self.finished { + return Poll::Ready(None); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + // Start the timer for compute time within this operator. It will be + // stopped when dropped. + let _timer = elapsed_compute.timer(); + + if self.search_mode.is_mode_linear() && batch.num_rows() > 0 { + self.most_recent_row = Some(get_last_row_batch(&batch)?); + } + self.search_mode.update_partition_batch( + &mut self.input_buffer, + batch, + &self.window_expr, + &mut self.partition_buffers, + )?; + if let Some(batch) = self.compute_aggregates()? { + return Poll::Ready(Some(Ok(batch))); + } + self.poll_next_inner(cx) + } + Some(Err(e)) => Poll::Ready(Some(Err(e))), + None => { + let _timer = elapsed_compute.timer(); + + self.finished = true; + // Release the input pipeline's resources before computing the + // final aggregates. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + for (_, partition_batch_state) in self.partition_buffers.iter_mut() { + partition_batch_state.is_end = true; + } + if let Some(batch) = self.compute_aggregates()? { + return Poll::Ready(Some(Ok(batch))); + } + Poll::Ready(None) + } + } + } + + /// Removes partitions that have ended. For the remaining partitions, + /// drops buffered rows that no window expression will need again. + fn prune_partition_batches(&mut self) { + // Check that per-state and per-partition end-flags are consistent; + // otherwise, the pruning code below might produce inconsistent state. + #[cfg(debug_assertions)] + for window_agg_state in self.window_agg_states.iter() { + for (partition_row, WindowState { state, .. }) in window_agg_state.iter() { + debug_assert_eq!( + state.is_end, self.partition_buffers[partition_row].is_end, + "window state's recorded end flag is out of sync with its partition" + ); + } + } + + // Remove partitions which we know already ended (is_end flag is true). + // Since the retain method preserves insertion order, we still have + // ordering in between partitions after removal. + self.partition_buffers + .retain(|_, partition_batch_state| !partition_batch_state.is_end); + // Likewise, drop per-window-expression state for ended partitions. + for window_agg_state in self.window_agg_states.iter_mut() { + window_agg_state.retain(|_, WindowState { state, .. }| !state.is_end); + } + + // Calculate how many rows to prune from each partition's batch. For a + // single window expression, rows before min(window_frame_range.start, + // last_calculated_index) are prunable: their results are already + // calculated, and frame boundaries never move backwards, so no future + // frame can include them. All window expressions share the partition + // batch, so a row can only be pruned once every expression is done with + // it: the count to prune is the minimum across expressions. A partition + // missing from the map has nothing to prune. + let mut n_prune_each_partition = HashMap::new(); + if let Some((first, rest)) = self.window_agg_states.split_first() { + // First window expression seeds the prune-count map + for (partition_row, WindowState { state, .. }) in first.iter() { + let n_prune = + min(state.window_frame_range.start, state.last_calculated_index); + if n_prune > 0 { + n_prune_each_partition.insert(partition_row.clone(), n_prune); + } + } + // Take the per-partition min of the prune-count for each + // additional window expression + for window_agg_state in rest { + n_prune_each_partition.retain(|partition_row, current| { + let Some(WindowState { state, .. }) = + window_agg_state.get(partition_row) + else { + return false; + }; + let n_prune = + min(state.window_frame_range.start, state.last_calculated_index); + *current = min(*current, n_prune); + *current > 0 + }); + } + } + + // Drop the prunable prefix of each partition's buffered batch: + for (partition_row, n_prune) in n_prune_each_partition.iter() { + debug_assert!( + *n_prune > 0, + "prune-count map must only contain positive entries" + ); + let pb_state = &mut self.partition_buffers[partition_row]; + + let batch = &pb_state.record_batch; + pb_state.record_batch = batch.slice(*n_prune, batch.num_rows() - n_prune); + + // Update state indices since we have pruned some rows from the beginning: + for window_agg_state in self.window_agg_states.iter_mut() { + window_agg_state[partition_row].state.prune_state(*n_prune); + } + } + } + + /// Prunes the section of the input batch whose aggregate results + /// are calculated and emitted. + fn prune_input_batch(&mut self, n_out: usize) -> Result<()> { + // Prune first n_out rows from the input_buffer + let n_to_keep = self.input_buffer.num_rows() - n_out; + let batch_to_keep = self + .input_buffer + .columns() + .iter() + .map(|elem| elem.slice(n_out, n_to_keep)) + .collect::<Vec<_>>(); + self.input_buffer = RecordBatch::try_new_with_options( + self.input_buffer.schema(), + batch_to_keep, + &RecordBatchOptions::new().with_row_count(Some(n_to_keep)), + )?; + Ok(()) + } + + /// Prunes emitted parts from WindowAggState `out_col` field. + fn prune_out_columns(&mut self) { + // We store generated columns for each window expression in the `out_col` + // field of `WindowAggState`. Given how many rows are emitted, we remove + // these sections from state. + for partition_window_agg_states in self.window_agg_states.iter_mut() { + // If `is_end` is set, directly remove the entry; this shrinks the + // hash map. + partition_window_agg_states + .retain(|_, partition_batch_state| !partition_batch_state.state.is_end); + } + // Only partitions that emitted rows since the previous pruning pass + // have output columns to shrink. Their emitted-row counts are + // consumed and reset here, so partitions that emitted nothing keep + // a count of zero and are passed over without any hash lookups. + for (partition_key, partition_batch) in self.partition_buffers.iter_mut() { + let n_emitted = partition_batch.n_out_row; + if n_emitted == 0 { + continue; + } + partition_batch.n_out_row = 0; + for partition_window_agg_states in self.window_agg_states.iter_mut() { + if let Some(WindowState { state, .. }) = + partition_window_agg_states.get_mut(partition_key) + { + let out_col = &mut state.out_col; + let n_to_keep = out_col.len() - n_emitted; + *out_col = out_col.slice(n_emitted, n_to_keep); + } + } + } + } +} + +impl RecordBatchStream for BoundedWindowAggStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +// Gets the index of minimum entry, returns None if empty. +fn argmin<T: PartialOrd>(data: impl Iterator<Item = T>) -> Option<(usize, T)> { + data.enumerate() + .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(Ordering::Equal)) +} + +/// Calculates the section we can show results for expression +fn get_aggregate_result_out_column( + partition_window_agg_states: &PartitionWindowAggStates, + len_to_show: usize, +) -> Result<ArrayRef> { + let mut result = None; + let mut running_length = 0; + let mut batches_to_concat = vec![]; + // We assume that iteration order is according to insertion order + for ( + _, + WindowState { + state: WindowAggState { out_col, .. }, + .. + }, + ) in partition_window_agg_states + { + if running_length < len_to_show { + let n_to_use = min(len_to_show - running_length, out_col.len()); + let slice_to_use = if n_to_use == out_col.len() { + // avoid slice when the entire column is used + Arc::clone(out_col) + } else { + out_col.slice(0, n_to_use) + }; + batches_to_concat.push(slice_to_use); + running_length += n_to_use; + } else { + break; + } + } + + if !batches_to_concat.is_empty() { + let array_refs: Vec<&dyn Array> = + batches_to_concat.iter().map(|a| a.as_ref()).collect(); + result = Some(concat(&array_refs)?); + } + + if running_length != len_to_show { + return exec_err!( + "Generated row number should be {len_to_show}, it is {running_length}" + ); + } + result.ok_or_else(|| exec_datafusion_err!("Should contain something")) +} + +/// Constructs a batch from the last row of batch in the argument. +pub(crate) fn get_last_row_batch(batch: &RecordBatch) -> Result<RecordBatch> { + if batch.num_rows() == 0 { + return exec_err!("Latest batch should have at least 1 row"); + } + Ok(batch.slice(batch.num_rows() - 1, 1)) +} + +#[cfg(test)] +mod tests { + use std::pin::Pin; + use std::sync::Arc; + use std::task::{Context, Poll}; + use std::time::Duration; + + use crate::common::collect; + use crate::execution_plan::CardinalityEffect; + use crate::expressions::PhysicalSortExpr; + use crate::projection::{ProjectionExec, ProjectionExpr}; + use crate::streaming::{PartitionStream, StreamingTableExec}; + use crate::test::TestMemoryExec; + use crate::windows::bounded_window_agg_exec::WindowStateObserver; + use crate::windows::{ + BoundedWindowAggExec, InputOrderMode, create_udwf_window_expr, create_window_expr, + }; + use crate::{ExecutionPlan, WindowExpr, displayable, execute_stream}; + + use arrow::array::{ + RecordBatch, + builder::{Int64Builder, UInt64Builder}, + }; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::test_util::batches_to_string; + use datafusion_common::{Result, ScalarValue, exec_datafusion_err}; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::{ + RecordBatchStream, SendableRecordBatchStream, TaskContext, + }; + use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, + }; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_functions_window::nth_value::last_value_udwf; + use datafusion_functions_window::nth_value::nth_value_udwf; + use datafusion_physical_expr::expressions::{Column, Literal, col}; + use datafusion_physical_expr::window::{PartitionKey, StandardWindowExpr}; + use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; + + use futures::future::Shared; + use futures::{FutureExt, Stream, StreamExt, pin_mut, ready}; + use insta::assert_snapshot; + use itertools::Itertools; + use tokio::time::timeout; + + #[derive(Debug, Clone)] + struct TestStreamPartition { + schema: SchemaRef, + batches: Vec<RecordBatch>, + idx: usize, + state: PolingState, + sleep_duration: Duration, + send_exit: bool, + } + + impl PartitionStream for TestStreamPartition { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream { + // We create an iterator from the record batches and map them into Ok values, + // converting the iterator into a futures::stream::Stream + Box::pin(self.clone()) + } + } + + impl Stream for TestStreamPartition { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + self.poll_next_inner(cx) + } + } + + #[derive(Debug, Clone)] + enum PolingState { + Sleep(Shared<futures::future::BoxFuture<'static, ()>>), + BatchReturn, + } + + impl TestStreamPartition { + fn poll_next_inner( + self: &mut Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + loop { + match &mut self.state { + PolingState::BatchReturn => { + // Wait for self.sleep_duration before sending any new data + let f = tokio::time::sleep(self.sleep_duration).boxed().shared(); + self.state = PolingState::Sleep(f); + let input_batch = if let Some(batch) = + self.batches.clone().get(self.idx) + { + batch.clone() + } else if self.send_exit { + // Send None to signal end of data + return Poll::Ready(None); + } else { + // Go to sleep mode + let f = + tokio::time::sleep(self.sleep_duration).boxed().shared(); + self.state = PolingState::Sleep(f); + continue; + }; + self.idx += 1; + return Poll::Ready(Some(Ok(input_batch))); + } + PolingState::Sleep(future) => { + pin_mut!(future); + ready!(future.poll_unpin(cx)); + self.state = PolingState::BatchReturn; + } + } + } + } + } + + impl RecordBatchStream for TestStreamPartition { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + } + + fn bounded_window_exec_pb_latent_range( + input: Arc<dyn ExecutionPlan>, + n_future_range: usize, + hash: &str, + order_by: &str, + ) -> Result<Arc<dyn ExecutionPlan>> { + let schema = input.schema(); + let window_fn = WindowFunctionDefinition::AggregateUDF(count_udaf()); + let col_expr = + Arc::new(Column::new(schema.fields[0].name(), 0)) as Arc<dyn PhysicalExpr>; + let args = vec![col_expr]; + let partitionby_exprs = vec![col(hash, &schema)?]; + let orderby_exprs = vec![PhysicalSortExpr { + expr: col(order_by, &schema)?, + options: SortOptions::default(), + }]; + let window_frame = WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(Some(n_future_range as u64))), + ); + let fn_name = format!( + "{window_fn}({args:?}) PARTITION BY: [{partitionby_exprs:?}], ORDER BY: [{orderby_exprs:?}]" + ); + let input_order_mode = InputOrderMode::Linear; + Ok(Arc::new(BoundedWindowAggExec::try_new( + vec![create_window_expr( + &window_fn, + fn_name, + &args, + &partitionby_exprs, + &orderby_exprs, + Arc::new(window_frame), + input.schema(), + false, + false, + None, + )?], + input, + input_order_mode, + true, + )?)) + } + + fn projection_exec(input: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn ExecutionPlan>> { + let schema = input.schema(); + let exprs = input + .schema() + .fields + .iter() + .enumerate() + .map(|(idx, field)| { + let name = if field.name().len() > 20 { + format!("col_{idx}") + } else { + field.name().clone() + }; + let expr = col(field.name(), &schema).unwrap(); + (expr, name) + }) + .collect::<Vec<_>>(); + let proj_exprs: Vec<ProjectionExpr> = exprs + .into_iter() + .map(|(expr, alias)| ProjectionExpr { expr, alias }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) + } + + fn task_context_helper() -> TaskContext { + let task_ctx = TaskContext::default(); + // Create session context with config + let session_config = SessionConfig::new() + .with_batch_size(1) + .with_target_partitions(2) + .with_round_robin_repartition(false); + task_ctx.with_session_config(session_config) + } + + fn task_context() -> Arc<TaskContext> { + Arc::new(task_context_helper()) + } + + pub async fn collect_stream( + mut stream: SendableRecordBatchStream, + results: &mut Vec<RecordBatch>, + ) -> Result<()> { + while let Some(item) = stream.next().await { + results.push(item?); + } + Ok(()) + } + + /// Execute the [ExecutionPlan] and collect the results in memory + pub async fn collect_with_timeout( + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, + timeout_duration: Duration, + ) -> Result<Vec<RecordBatch>> { + let stream = execute_stream(plan, context)?; + let mut results = vec![]; + + // Execute the asynchronous operation with a timeout + if timeout(timeout_duration, collect_stream(stream, &mut results)) + .await + .is_ok() + { + return Err(exec_datafusion_err!("shouldn't have completed")); + }; + + Ok(results) + } + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("sn", DataType::UInt64, true), + Field::new("hash", DataType::Int64, true), + ])) + } + + fn schema_orders(schema: &SchemaRef) -> Result<Vec<LexOrdering>> { + let orderings = vec![ + [PhysicalSortExpr { + expr: col("sn", schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }] + .into(), + ]; + Ok(orderings) + } + + fn is_integer_division_safe(lhs: usize, rhs: usize) -> bool { + let res = lhs / rhs; + res * rhs == lhs + } + fn generate_batches( + schema: &SchemaRef, + n_row: usize, + n_chunk: usize, + ) -> Result<Vec<RecordBatch>> { + let mut batches = vec![]; + assert!(n_row > 0); + assert!(n_chunk > 0); + assert!(is_integer_division_safe(n_row, n_chunk)); + let hash_replicate = 4; + + let chunks = (0..n_row) + .chunks(n_chunk) + .into_iter() + .map(|elem| elem.into_iter().collect::<Vec<_>>()) + .collect::<Vec<_>>(); + + // Send 2 RecordBatches at the source + for sn_values in chunks { + let mut sn1_array = UInt64Builder::with_capacity(sn_values.len()); + let mut hash_array = Int64Builder::with_capacity(sn_values.len()); + + for sn in sn_values { + sn1_array.append_value(sn as u64); + let hash_value = (2 - (sn / hash_replicate)) as i64; + hash_array.append_value(hash_value); + } + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(sn1_array.finish()), Arc::new(hash_array.finish())], + )?; + batches.push(batch); + } + Ok(batches) + } + + fn generate_never_ending_source( + n_rows: usize, + chunk_length: usize, + n_partition: usize, + is_infinite: bool, + send_exit: bool, + per_batch_wait_duration_in_millis: u64, + ) -> Result<Arc<dyn ExecutionPlan>> { + assert!(n_partition > 0); + + // We use same hash value in the table. This makes sure that + // After hashing computation will continue in only in one of the output partitions + // In this case, data flow should still continue + let schema = test_schema(); + let orderings = schema_orders(&schema)?; + + // Source waits per_batch_wait_duration_in_millis ms before sending other batch + let per_batch_wait_duration = + Duration::from_millis(per_batch_wait_duration_in_millis); + + let batches = generate_batches(&schema, n_rows, chunk_length)?; + + // Source has 2 partitions + let partitions = vec![ + Arc::new(TestStreamPartition { + schema: Arc::clone(&schema), + batches, + idx: 0, + state: PolingState::BatchReturn, + sleep_duration: per_batch_wait_duration, + send_exit, + }) as _; + n_partition + ]; + let source = Arc::new(StreamingTableExec::try_new( + Arc::clone(&schema), + partitions, + None, + orderings, + is_infinite, + None, + )?) as _; + Ok(source) + } + + // Tests NTH_VALUE(negative index) with memoize feature + // To be able to trigger memoize feature for NTH_VALUE we need to + // - feed BoundedWindowAggExec with batch stream data. + // - Window frame should contain UNBOUNDED PRECEDING. + // It hard to ensure these conditions are met, from the sql query. + #[tokio::test] + async fn test_window_nth_value_bounded_memoize() -> Result<()> { + let config = SessionConfig::new().with_target_partitions(1); + let task_ctx = Arc::new(TaskContext::default().with_session_config(config)); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + // Create a new batch of data to insert into the table + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3]))], + )?; + + let memory_exec = TestMemoryExec::try_new_exec( + &[vec![batch.clone(), batch.clone(), batch.clone()]], + Arc::clone(&schema), + None, + )?; + let col_a = col("a", &schema)?; + let nth_value_func1 = create_udwf_window_expr( + &nth_value_udwf(), + &[ + Arc::clone(&col_a), + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + ], + &schema, + "nth_value(-1)".to_string(), + false, + )? + .reverse_expr() + .unwrap(); + let nth_value_func2 = create_udwf_window_expr( + &nth_value_udwf(), + &[ + Arc::clone(&col_a), + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + ], + &schema, + "nth_value(-2)".to_string(), + false, + )? + .reverse_expr() + .unwrap(); + + let last_value_func = create_udwf_window_expr( + &last_value_udwf(), + &[Arc::clone(&col_a)], + &schema, + "last".to_string(), + false, + )?; + + let window_exprs = vec![ + // LAST_VALUE(a) + Arc::new(StandardWindowExpr::new( + last_value_func, + &[], + &[], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + )) as _, + // NTH_VALUE(a, -1) + Arc::new(StandardWindowExpr::new( + nth_value_func1, + &[], + &[], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + )) as _, + // NTH_VALUE(a, -2) + Arc::new(StandardWindowExpr::new( + nth_value_func2, + &[], + &[], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + )) as _, + ]; + let physical_plan = BoundedWindowAggExec::try_new( + window_exprs, + memory_exec, + InputOrderMode::Sorted, + true, + ) + .map(|e| Arc::new(e) as Arc<dyn ExecutionPlan>)?; + + let batches = collect(physical_plan.execute(0, task_ctx)?).await?; + + // Get string representation of the plan + assert_snapshot!(displayable(physical_plan.as_ref()).indent(true), @r#" + BoundedWindowAggExec: wdw=[last: Field { "last": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, nth_value(-1): Field { "nth_value(-1)": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, nth_value(-2): Field { "nth_value(-2)": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + DataSourceExec: partitions=1, partition_sizes=[3] + "#); + + assert_snapshot!(batches_to_string(&batches), @r" + +---+------+---------------+---------------+ + | a | last | nth_value(-1) | nth_value(-2) | + +---+------+---------------+---------------+ + | 1 | 1 | 1 | | + | 2 | 2 | 2 | 1 | + | 3 | 3 | 3 | 2 | + | 1 | 1 | 1 | 3 | + | 2 | 2 | 2 | 1 | + | 3 | 3 | 3 | 2 | + | 1 | 1 | 1 | 3 | + | 2 | 2 | 2 | 1 | + | 3 | 3 | 3 | 2 | + +---+------+---------------+---------------+ + "); + Ok(()) + } + + // In `Linear` mode, a partition may receive no new rows for several + // input batches while other partitions keep growing. Once all of a + // partition's buffered rows have results, the evaluation sweep skips + // it until it receives rows again, so this test drives a partition + // through quiet batches and then resumes it: the results after the + // gap must continue from the retained accumulator state. Both frames + // are causal, so results finalize in the batch their row arrives in + // and the quiet partition is fully calculated while it waits. + #[tokio::test] + async fn bounded_window_linear_quiet_partition_resume() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::UInt64, false), + Field::new("ts", DataType::UInt64, false), + ])); + let make_batch = |rows: &[(u64, u64)]| -> Result<RecordBatch> { + let mut pk = UInt64Builder::with_capacity(rows.len()); + let mut ts = UInt64Builder::with_capacity(rows.len()); + for (p, t) in rows { + pk.append_value(*p); + ts.append_value(*t); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(pk.finish()), Arc::new(ts.finish())], + )?) + }; + // `ts` ascends globally; partition 0 is absent from the middle batches. + let batches = vec![ + make_batch(&[(0, 0), (0, 1), (1, 2)])?, + make_batch(&[(1, 3), (1, 4)])?, + make_batch(&[(1, 5)])?, + make_batch(&[(0, 6), (1, 7)])?, + ]; + let memory_exec = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let partition_by = vec![col("pk", &schema)?]; + let order_by = [PhysicalSortExpr { + expr: col("ts", &schema)?, + options: SortOptions::default(), + }]; + // A running COUNT (plain aggregate) and a SUM over the previous and + // current row (sliding aggregate). + let count_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let sum_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let physical_plan = BoundedWindowAggExec::try_new( + vec![count_expr, sum_expr], + memory_exec, + InputOrderMode::Linear, + true, + ) + .map(|e| Arc::new(e) as Arc<dyn ExecutionPlan>)?; + + let batches = collect(physical_plan.execute(0, task_context())?).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-------+-----+ + | pk | ts | count | sum | + +----+----+-------+-----+ + | 0 | 0 | 1 | 0 | + | 0 | 1 | 2 | 1 | + | 1 | 2 | 1 | 2 | + | 1 | 3 | 2 | 5 | + | 1 | 4 | 3 | 7 | + | 1 | 5 | 4 | 9 | + | 0 | 6 | 3 | 7 | + | 1 | 7 | 5 | 12 | + +----+----+-------+-----+ + "); + Ok(()) + } + + // This test, tests whether most recent row guarantee by the input batch of the `BoundedWindowAggExec` + // helps `BoundedWindowAggExec` to generate low latency result in the `Linear` mode. + // Input data generated at the source is + // "+----+------+", + // "| sn | hash |", + // "+----+------+", + // "| 0 | 2 |", + // "| 1 | 2 |", + // "| 2 | 2 |", + // "| 3 | 2 |", + // "| 4 | 1 |", + // "| 5 | 1 |", + // "| 6 | 1 |", + // "| 7 | 1 |", + // "| 8 | 0 |", + // "| 9 | 0 |", + // "+----+------+", + // + // Effectively following query is run on this data + // + // SELECT *, count(*) OVER(PARTITION BY duplicated_hash ORDER BY sn RANGE BETWEEN CURRENT ROW AND 1 FOLLOWING) + // FROM test; + // + // partition `duplicated_hash=2` receives following data from the input + // + // "+----+------+", + // "| sn | hash |", + // "+----+------+", + // "| 0 | 2 |", + // "| 1 | 2 |", + // "| 2 | 2 |", + // "| 3 | 2 |", + // "+----+------+", + // normally `BoundedWindowExec` can only generate following result from the input above + // + // "+----+------+---------+", + // "| sn | hash | count |", + // "+----+------+---------+", + // "| 0 | 2 | 2 |", + // "| 1 | 2 | 2 |", + // "| 2 | 2 |<not yet>|", + // "| 3 | 2 |<not yet>|", + // "+----+------+---------+", + // where result of last 2 row is missing. Since window frame end is not may change with future data + // since window frame end is determined by 1 following (To generate result for row=3[where sn=2] we + // need to received sn=4 to make sure window frame end bound won't change with future data). + // + // With the ability of different partitions to use global ordering at the input (where most up-to date + // row is + // "| 9 | 0 |", + // ) + // + // `BoundedWindowExec` should be able to generate following result in the test + // + // "+----+------+-------+", + // "| sn | hash | col_2 |", + // "+----+------+-------+", + // "| 0 | 2 | 2 |", + // "| 1 | 2 | 2 |", + // "| 2 | 2 | 2 |", + // "| 3 | 2 | 1 |", + // "| 4 | 1 | 2 |", + // "| 5 | 1 | 2 |", + // "| 6 | 1 | 2 |", + // "| 7 | 1 | 1 |", + // "+----+------+-------+", + // + // where result for all rows except last 2 is calculated (To calculate result for row 9 where sn=8 + // we need to receive sn=10 value to calculate it result.). + // In this test, out aim is to test for which portion of the input data `BoundedWindowExec` can generate + // a result. To test this behaviour, we generated the data at the source infinitely (no `None` signal + // is sent to output from source). After, row: + // + // "| 9 | 0 |", + // + // is sent. Source stops sending data to output. We collect, result emitted by the `BoundedWindowExec` at the + // end of the pipeline with a timeout (Since no `None` is sent from source. Collection never ends otherwise). + #[tokio::test] + async fn bounded_window_exec_linear_mode_range_information() -> Result<()> { + let n_rows = 10; + let chunk_length = 2; + let n_future_range = 1; + + let timeout_duration = Duration::from_millis(2000); + + let source = + generate_never_ending_source(n_rows, chunk_length, 1, true, false, 5)?; + + let window = + bounded_window_exec_pb_latent_range(source, n_future_range, "hash", "sn")?; + + let plan = projection_exec(window)?; + + // Get string representation of the plan + assert_snapshot!(displayable(plan.as_ref()).indent(true), @r#" + ProjectionExec: expr=[sn@0 as sn, hash@1 as hash, count([Column { name: "sn", index: 0 }]) PARTITION BY: [[Column { name: "hash", index: 1 }]], ORDER BY: [[PhysicalSortExpr { expr: Column { name: "sn", index: 0 }, options: SortOptions { descending: false, nulls_first: true } }]]@2 as col_2] + BoundedWindowAggExec: wdw=[count([Column { name: "sn", index: 0 }]) PARTITION BY: [[Column { name: "hash", index: 1 }]], ORDER BY: [[PhysicalSortExpr { expr: Column { name: "sn", index: 0 }, options: SortOptions { descending: false, nulls_first: true } }]]: Field { "count([Column { name: \"sn\", index: 0 }]) PARTITION BY: [[Column { name: \"hash\", index: 1 }]], ORDER BY: [[PhysicalSortExpr { expr: Column { name: \"sn\", index: 0 }, options: SortOptions { descending: false, nulls_first: true } }]]": Int64 }, frame: RANGE BETWEEN CURRENT ROW AND 1 FOLLOWING], mode=[Linear] + StreamingTableExec: partition_sizes=1, projection=[sn, hash], infinite_source=true, output_ordering=[sn@0 ASC NULLS LAST] + "#); + + let task_ctx = task_context(); + let batches = collect_with_timeout(plan, task_ctx, timeout_duration).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+------+-------+ + | sn | hash | col_2 | + +----+------+-------+ + | 0 | 2 | 2 | + | 1 | 2 | 2 | + | 2 | 2 | 2 | + | 3 | 2 | 1 | + | 4 | 1 | 2 | + | 5 | 1 | 2 | + | 6 | 1 | 2 | + | 7 | 1 | 1 | + +----+------+-------+ + "); + + Ok(()) + } + + type Observation = (usize, PartitionKey, Vec<ScalarValue>); + + /// Test [`WindowStateObserver`] that records every callback into a shared + /// `Vec` for later assertion. + struct RecordingObserver { + sink: Arc<std::sync::Mutex<Vec<Observation>>>, + } + + impl WindowStateObserver for RecordingObserver { + fn finalize_window_aggregate( + &self, + partition_idx: usize, + _window_expr: &Arc<dyn WindowExpr>, + partition_key: &PartitionKey, + state: Vec<ScalarValue>, + ) -> Result<()> { + self.sink + .lock() + .unwrap() + .push((partition_idx, partition_key.clone(), state)); + Ok(()) + } + } + + /// Build a `BoundedWindowAggExec` for `count(sn) OVER (PARTITION BY hash + /// ORDER BY sn <frame>)` over a fixed two-group source (hash=1 × 3, + /// hash=2 × 3, sorted by (hash, sn)). Returns the plan pre-observer so + /// callers can decide how to install it. + fn build_partition_close_plan(frame: WindowFrame) -> Result<BoundedWindowAggExec> { + let schema = test_schema(); + + let mut sn_b = UInt64Builder::with_capacity(6); + let mut hash_b = Int64Builder::with_capacity(6); + for (sn, hash) in [(1u64, 1i64), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2)] { + sn_b.append_value(sn); + hash_b.append_value(hash); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [ + PhysicalSortExpr { + expr: col("hash", &schema)?, + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }, + ] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "cnt".to_string(), + &[col("sn", &schema)?], + &[col("hash", &schema)?], + &[PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }], + Arc::new(frame), + source.schema(), + false, + false, + None, + )?; + + BoundedWindowAggExec::try_new(vec![expr], source, InputOrderMode::Sorted, false) + } + + // Two PARTITION BY groups: hash=1 [sn=1,2,3] then hash=2 [sn=4,5,6]. + // Input is sorted by (hash, sn) so we can run in Sorted mode; in that + // mode `mark_partition_end` closes the leading group mid-stream and + // EOS closes the tail — both fire the observer for an ever-expanding + // frame. Sliding frames are rejected at install time. + + #[tokio::test] + async fn test_state_observer_rejects_sliding_frame() -> Result<()> { + // `CURRENT ROW → UNBOUNDED FOLLOWING` is not ever-expanding, so this + // maps to `SlidingAggregateWindowExpr` whose accumulator retracts as + // rows leave the frame — at partition close the accumulator holds + // only the last frame's rows, not the partition aggregate. + // `with_state_observer` refuses this configuration. + use std::sync::Mutex; + + let plan = build_partition_close_plan(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(None)), + ))?; + let observer: Arc<dyn WindowStateObserver> = Arc::new(RecordingObserver { + sink: Arc::new(Mutex::new(vec![])), + }); + let err = plan.with_state_observer(Some(observer)).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("sliding aggregate window frame"), + "expected sliding-frame rejection, got: {msg}" + ); + Ok(()) + } + + #[tokio::test] + async fn test_finalized_state_observer_fires_on_causal_frame() -> Result<()> { + // `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` — ever-expanding, + // `PlainAggregateWindowExpr` under the hood. At partition close the + // accumulator holds the partition aggregate. Both mid-stream close + // (hash=1 as hash=2 rows arrive) and EOS (hash=2 at drain) fire. + use std::sync::Mutex; + + let task_ctx = Arc::new(TaskContext::default()); + let plan = build_partition_close_plan(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + ))?; + + let observations: Arc<Mutex<Vec<Observation>>> = Arc::new(Mutex::new(vec![])); + let observer: Arc<dyn WindowStateObserver> = Arc::new(RecordingObserver { + sink: Arc::clone(&observations), + }); + let plan = plan.with_state_observer(Some(observer))?; + + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + // count(sn) over each of hash=1 (3 rows) and hash=2 (3 rows), in + // close order — hash=1 first (mid-stream close), hash=2 second (EOS). + let observed: Vec<(usize, i64, Vec<ScalarValue>)> = observations + .lock() + .unwrap() + .iter() + .map(|(idx, key, state)| { + let hash = match &key[0] { + ScalarValue::Int64(Some(v)) => *v, + other => panic!("unexpected partition-key element: {other:?}"), + }; + (*idx, hash, state.clone()) + }) + .collect(); + assert_eq!( + observed, + vec![ + (0, 1, vec![ScalarValue::Int64(Some(3))]), + (0, 2, vec![ScalarValue::Int64(Some(3))]), + ] + ); + Ok(()) + } + + #[tokio::test] + async fn test_finalized_state_observer_fires_exactly_once_across_batches() + -> Result<()> { + // Regression guard for the exactly-once observer contract when + // partition close and pruning happen on different `compute_aggregates` + // calls. + // + // The observer fires from `publish_finalized_states`, called at the + // top of every `compute_aggregates`. Entries are only cleared by + // `prune_state`, which runs only when `calculate_out_columns` returns + // `Some`. Nothing in the type system ties the two together, so a + // group whose state was published on batch N must not be re-published + // on batch N+1 or at EOS. + // + // Layout: three PARTITION BY groups streamed across two input + // batches, so each group closes on a distinct `compute_aggregates` + // call: + // batch 1 = [hash=1 × 2] — no close (single group). + // batch 2 = [hash=2 × 2, hash=3 × 2] — `mark_partition_end` + // closes hash=1 and hash=2. + // EOS — closes hash=3. + // + // Assertion: each key appears exactly once across all observations. + use std::sync::Mutex; + + let task_ctx = Arc::new(TaskContext::default()); + let schema = test_schema(); + + // Two batches, same output partition. + let make_batch = |rows: &[(u64, i64)]| -> Result<RecordBatch> { + let mut sn_b = UInt64Builder::with_capacity(rows.len()); + let mut hash_b = Int64Builder::with_capacity(rows.len()); + for &(sn, hash) in rows { + sn_b.append_value(sn); + hash_b.append_value(hash); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?) + }; + let batch1 = make_batch(&[(1, 1), (2, 1)])?; + let batch2 = make_batch(&[(3, 2), (4, 2), (5, 3), (6, 3)])?; + + let ordering: LexOrdering = [ + PhysicalSortExpr { + expr: col("hash", &schema)?, + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }, + ] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch1, batch2]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "cnt".to_string(), + &[col("sn", &schema)?], + &[col("hash", &schema)?], + &[PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + source.schema(), + false, + false, + None, + )?; + + let observations: Arc<Mutex<Vec<Observation>>> = Arc::new(Mutex::new(vec![])); + let observer: Arc<dyn WindowStateObserver> = Arc::new(RecordingObserver { + sink: Arc::clone(&observations), + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_state_observer(Some(observer))?; + + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + let fired: Vec<i64> = observations + .lock() + .unwrap() + .iter() + .map(|(_, key, _)| match &key[0] { + ScalarValue::Int64(Some(v)) => *v, + other => panic!("unexpected partition-key element: {other:?}"), + }) + .collect(); + // Each group closes on a distinct `compute_aggregates` call — hash=1 + // and hash=2 on batch 2's `mark_partition_end`, hash=3 at EOS — and + // each appears exactly once, in close order. + assert_eq!(fired, vec![1, 2, 3]); + Ok(()) + } + + /// Run one task's local BWAG for `SUM(sn) OVER (ORDER BY sn ROWS + /// UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, over + /// `input` sorted ascending. Returns the per-row output values and the + /// observed finalized state total (which the caller uses as a carry-in + /// for the next task). + async fn run_running_sum_task( + input: &[u64], + task_ctx: Arc<TaskContext>, + ) -> Result<(Vec<u64>, u64)> { + use arrow::array::UInt64Array; + use datafusion_functions_aggregate::sum::sum_udaf; + use std::sync::Mutex; + + /// Observer for `run_running_sum_task`: captures the single running + /// SUM total published at EOS. Asserts exactly-one fire and rejects + /// non-empty partition keys (this helper is no-PARTITION-BY only). + struct RunningSumObserver { + sink: Arc<Mutex<Option<u64>>>, + } + + impl WindowStateObserver for RunningSumObserver { + fn finalize_window_aggregate( + &self, + _partition_idx: usize, + _window_expr: &Arc<dyn WindowExpr>, + partition_key: &PartitionKey, + state: Vec<ScalarValue>, + ) -> Result<()> { + assert!( + partition_key.is_empty(), + "empty PartitionKey for no-PARTITION-BY plan" + ); + let total = match &state[0] { + ScalarValue::UInt64(Some(v)) => *v, + ScalarValue::Int64(Some(v)) => *v as u64, + other => panic!("unexpected sum state element: {other:?}"), + }; + let prev = self.sink.lock().unwrap().replace(total); + assert!(prev.is_none(), "observer must fire exactly once per task"); + Ok(()) + } + } + + let schema = test_schema(); + let mut sn_b = UInt64Builder::with_capacity(input.len()); + let mut hash_b = Int64Builder::with_capacity(input.len()); + for &sn in input { + sn_b.append_value(sn); + hash_b.append_value(0); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let window_fn = WindowFunctionDefinition::AggregateUDF(sum_udaf()); + let args = vec![col("sn", &schema)?]; + let partition_by: Vec<Arc<dyn PhysicalExpr>> = vec![]; + let order_by = vec![PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }]; + let frame = WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + ); + let expr = create_window_expr( + &window_fn, + "running_sum".to_string(), + &args, + &partition_by, + &order_by, + Arc::new(frame), + source.schema(), + false, + false, + None, + )?; + + let total_sink: Arc<Mutex<Option<u64>>> = Arc::new(Mutex::new(None)); + let observer: Arc<dyn WindowStateObserver> = Arc::new(RunningSumObserver { + sink: Arc::clone(&total_sink), + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_state_observer(Some(observer))?; + let batches = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + let mut out = Vec::with_capacity(input.len()); + for batch in &batches { + let col = batch + .column_by_name("running_sum") + .expect("running_sum column present"); + let arr = col + .as_any() + .downcast_ref::<UInt64Array>() + .expect("SUM(UInt64) → UInt64Array"); + for i in 0..arr.len() { + out.push(arr.value(i)); + } + } + let total = total_sink + .lock() + .unwrap() + .expect("observer must have fired at EOS"); + Ok((out, total)) + } + + /// Run one task's local BWAG for `approx_distinct(sn) OVER (ORDER BY sn + /// ROWS UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, and + /// return the single EOS-observed [`Accumulator::state`] Vec. + async fn run_approx_distinct_task( + input: &[u64], + task_ctx: Arc<TaskContext>, + ) -> Result<Vec<ScalarValue>> { + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; + use std::sync::Mutex; + + /// Observer for `run_approx_distinct_task`: capture the single EOS + /// state. Asserts exactly-one fire and rejects non-empty partition + /// keys (helper is no-PARTITION-BY only). + struct ApproxDistinctObserver { + sink: Arc<Mutex<Option<Vec<ScalarValue>>>>, + } + + impl WindowStateObserver for ApproxDistinctObserver { + fn finalize_window_aggregate( + &self, + _partition_idx: usize, + _window_expr: &Arc<dyn WindowExpr>, + partition_key: &PartitionKey, + state: Vec<ScalarValue>, + ) -> Result<()> { + assert!( + partition_key.is_empty(), + "empty PartitionKey for no-PARTITION-BY plan" + ); + let prev = self.sink.lock().unwrap().replace(state); + assert!(prev.is_none(), "observer must fire exactly once per task"); + Ok(()) + } + } + + let schema = test_schema(); + let mut sn_b = UInt64Builder::with_capacity(input.len()); + let mut hash_b = Int64Builder::with_capacity(input.len()); + for &sn in input { + sn_b.append_value(sn); + hash_b.append_value(0); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(approx_distinct_udaf()), + "approx_distinct_sn".to_string(), + &[col("sn", &schema)?], + &[], + &[PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + source.schema(), + false, + false, + None, + )?; + + let state_sink: Arc<Mutex<Option<Vec<ScalarValue>>>> = Arc::new(Mutex::new(None)); + let observer: Arc<dyn WindowStateObserver> = Arc::new(ApproxDistinctObserver { + sink: Arc::clone(&state_sink), + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_state_observer(Some(observer))?; + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + state_sink + .lock() + .unwrap() + .take() + .ok_or_else(|| exec_datafusion_err!("observer never fired")) + } + + #[tokio::test] + async fn test_prefix_scan_across_tasks_matches_single_bwag() -> Result<()> { + // Demonstrates the parallel-window shape reviewers asked about: + // range-shuffle `SUM(sn) OVER (ORDER BY sn UNBOUNDED PRECEDING TO + // CURRENT ROW)` across two tasks, then prefix-scan each task's + // finalized state (from the observer) to carry-in the next task's + // rows. Result must match a single BWAG over the concatenated input. + let task_ctx = Arc::new(TaskContext::default()); + + // Two tasks under range partition on sn: + let (task1_out, task1_total) = + run_running_sum_task(&[1, 1, 2, 2, 3, 3, 4, 4], Arc::clone(&task_ctx)) + .await?; + let (task2_out, task2_total) = + run_running_sum_task(&[5, 5, 6, 6, 7, 7, 8, 8], Arc::clone(&task_ctx)) + .await?; + + // Local (uncorrected) outputs and totals — first pass. + assert_eq!(task1_out, vec![1, 2, 4, 6, 9, 12, 16, 20]); + assert_eq!(task1_total, 20); + assert_eq!(task2_out, vec![5, 10, 16, 22, 29, 36, 44, 52]); + assert_eq!(task2_total, 52); + + // Prefix scan over per-task totals → carry-in for each task. Task 0's + // carry-in is 0; task N's carry-in is the sum of tasks [0, N). + let carry_ins = [0u64, task1_total]; + + // Second pass: shift each task's local values by its carry-in. + let task1_final: Vec<u64> = task1_out.iter().map(|v| v + carry_ins[0]).collect(); + let task2_final: Vec<u64> = task2_out.iter().map(|v| v + carry_ins[1]).collect(); + let parallel_result: Vec<u64> = task1_final + .iter() + .chain(task2_final.iter()) + .copied() + .collect(); + + // Oracle: single BWAG over the full concatenated input. + let (single_result, single_total) = run_running_sum_task( + &[1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8], + task_ctx, + ) + .await?; + + assert_eq!( + parallel_result, single_result, + "two-task prefix-scan must match single-BWAG oracle" + ); + // And matches the sequence in the design discussion. + assert_eq!( + single_result, + vec![1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, 42, 49, 56, 64, 72] + ); + assert_eq!(single_total, 72); + Ok(()) + } + + #[tokio::test] + async fn test_prefix_merge_across_tasks_approx_distinct() -> Result<()> { + // Load-bearing contract for the parallel-window use case: the state + // exposed by `WindowStateObserver::finalize_window_aggregate` must be + // compatible with `Accumulator::merge_batch` on a fresh accumulator + // of the same UDAF. This is what allows non-decomposable aggregates + // like `approx_distinct` (HLL sketch state) to be prefix-merged + // across shard tasks — the reason we exposed accumulator state at + // all. If this ever breaks, downstream parallel-window work has to + // wait for a public API change. + use arrow::array::{ArrayRef, BinaryArray}; + use arrow::datatypes::FieldRef; + use datafusion_expr::function::AccumulatorArgs; + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; + + let task_ctx = Arc::new(TaskContext::default()); + + // Two tasks with overlapping inputs; concatenated distinct universe + // is {1,2,3,4,5}. + let state1 = + run_approx_distinct_task(&[1, 1, 2, 3], Arc::clone(&task_ctx)).await?; + let state2 = run_approx_distinct_task(&[3, 4, 5], Arc::clone(&task_ctx)).await?; + let state_single = + run_approx_distinct_task(&[1, 1, 2, 3, 3, 4, 5], Arc::clone(&task_ctx)) + .await?; + + // approx_distinct state is a single serialized-HLL Binary field. + assert_eq!(state1.len(), 1, "single state field"); + assert_eq!(state2.len(), 1, "single state field"); + assert_eq!(state_single.len(), 1, "single state field"); + + // Seed a fresh accumulator with the given serialized HLL states via + // `merge_batch` and return its distinct-count evaluation. + fn evaluate_merged(states: &[&ScalarValue]) -> Result<ScalarValue> { + let udaf = approx_distinct_udaf(); + let input_schema = + Arc::new(Schema::new(vec![Field::new("sn", DataType::UInt64, true)])); + let return_field: FieldRef = + Arc::new(Field::new("approx_distinct_sn", DataType::UInt64, true)); + let expr_field: FieldRef = Arc::new(Field::new("sn", DataType::UInt64, true)); + let physical_col: Arc<dyn PhysicalExpr> = col("sn", &input_schema)?; + let args = AccumulatorArgs { + return_field: Arc::clone(&return_field), + schema: &input_schema, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "approx_distinct", + is_distinct: false, + exprs: std::slice::from_ref(&physical_col), + expr_fields: std::slice::from_ref(&expr_field), + }; + let mut acc = udaf.accumulator(args)?; + let byte_slices: Vec<&[u8]> = states + .iter() + .map(|s| match s { + ScalarValue::Binary(Some(v)) => v.as_slice(), + other => panic!("expected Binary state, got {other:?}"), + }) + .collect(); + let bin: ArrayRef = Arc::new(BinaryArray::from_iter_values(byte_slices)); + acc.merge_batch(std::slice::from_ref(&bin))?; + acc.evaluate() + } + + let merged = evaluate_merged(&[&state1[0], &state2[0]])?; + let oracle = evaluate_merged(&[&state_single[0]])?; + + assert_eq!( + merged, oracle, + "merged task states must match single-BWAG oracle — parallel prefix-merge contract" + ); + // HLL is approximate but exact for a 5-element universe. + assert_eq!(merged, ScalarValue::UInt64(Some(5))); + Ok(()) + } + + #[test] + fn test_bounded_window_agg_cardinality_effect() -> Result<()> { + let schema = test_schema(); + let input: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + let plan = bounded_window_exec_pb_latent_range(input, 1, "hash", "sn")?; + let plan = plan + .downcast_ref::<BoundedWindowAggExec>() + .expect("expected BoundedWindowAggExec"); + + assert!(matches!( + plan.cardinality_effect(), + CardinalityEffect::Equal + )); + Ok(()) + } + + /// Checks the per-partition batches that `LinearSearch` splits an input + /// batch into: partitions appear in first-appearance order, rows within a + /// partition keep their stream order, NULL keys form their own partition, + /// and a single-partition batch is passed through without copying. + #[test] + fn test_linear_search_evaluate_partition_batches() -> Result<()> { + use super::{LinearSearch, PartitionSearcher}; + use arrow::array::{Int32Array, Int64Array}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int64, false), + ])); + let window_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count".to_string(), + &[col("b", &schema)?], + &[col("a", &schema)?], + &[], + Arc::new(WindowFrame::new(None)), + Arc::clone(&schema), + false, + false, + None, + )?; + let mut searcher = LinearSearch::new(vec![], Arc::clone(&schema)); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(1), + None, + Some(2), + Some(1), + ])), + Arc::new(Int64Array::from(vec![10, 20, 11, 30, 21, 12])), + ], + )?; + let result = + searcher.evaluate_partition_batches(&batch, &[Arc::clone(&window_expr)])?; + assert_eq!(result.len(), 3); + let expected = [ + ( + ScalarValue::Int32(Some(1)), + vec![Some(1); 3], + vec![10i64, 11, 12], + ), + (ScalarValue::Int32(Some(2)), vec![Some(2); 2], vec![20, 21]), + (ScalarValue::Int32(None), vec![None], vec![30]), + ]; + for ((key, partition_batch), (exp_key, exp_a, exp_b)) in + result.iter().zip(expected) + { + assert_eq!(key, &vec![exp_key]); + let exp_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(exp_a)), + Arc::new(Int64Array::from(exp_b)), + ], + )?; + assert_eq!(partition_batch, &exp_batch); + } + + let single = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(7), Some(7)])), + Arc::new(Int64Array::from(vec![70, 71])), + ], + )?; + let result = searcher.evaluate_partition_batches(&single, &[window_expr])?; + assert_eq!(result.len(), 1); + assert_eq!(result[0].0, vec![ScalarValue::Int32(Some(7))]); + assert_eq!(result[0].1, single); + // The whole batch belongs to one partition, so its columns are reused + // rather than gathered into a new batch. + assert!(Arc::ptr_eq(result[0].1.column(0), single.column(0))); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/windows/mod.rs b/native/vendor/datafusion-physical-plan/src/windows/mod.rs new file mode 100644 index 0000000..089bdc2 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/windows/mod.rs
@@ -0,0 +1,1385 @@ +// 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. + +//! Physical expressions for window functions + +mod bounded_window_agg_exec; +#[cfg(feature = "proto")] +mod proto; +mod utils; +mod window_agg_exec; + +use std::borrow::Borrow; +use std::sync::Arc; + +use crate::{ + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PhysicalExpr, + expressions::PhysicalSortExpr, +}; + +use arrow::datatypes::{Schema, SchemaRef}; +use arrow_schema::{FieldRef, SortOptions}; +use datafusion_common::{Result, exec_err}; +use datafusion_expr::{ + LimitEffect, PartitionEvaluator, ReversedUDWF, SetMonotonicity, WindowFrame, + WindowFunctionDefinition, WindowUDF, +}; +use datafusion_functions_window_common::expr::ExpressionArgs; +use datafusion_functions_window_common::field::WindowUDFFieldArgs; +use datafusion_functions_window_common::partition::PartitionEvaluatorArgs; +use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::window::{ + SlidingAggregateWindowExpr, StandardWindowFunctionExpr, +}; +use datafusion_physical_expr::{ConstExpr, EquivalenceProperties}; +use datafusion_physical_expr_common::sort_expr::{ + LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortRequirement, +}; + +use itertools::Itertools; + +// Public interface: +pub use bounded_window_agg_exec::{BoundedWindowAggExec, WindowStateObserver}; +pub use datafusion_physical_expr::window::{ + PlainAggregateWindowExpr, StandardWindowExpr, WindowExpr, +}; +pub use window_agg_exec::WindowAggExec; + +/// Build field from window function and add it into schema +pub fn schema_add_window_field( + args: &[Arc<dyn PhysicalExpr>], + schema: &Schema, + window_fn: &WindowFunctionDefinition, + fn_name: &str, +) -> Result<Arc<Schema>> { + let fields = args + .iter() + .map(|e| Arc::clone(e).as_ref().return_field(schema)) + .collect::<Result<Vec<_>>>()?; + let window_expr_return_field = window_fn.return_field(&fields, fn_name)?; + let mut window_fields = schema + .fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect_vec(); + // Skip extending schema for UDAF + if let WindowFunctionDefinition::AggregateUDF(_) = window_fn { + Ok(Arc::new(Schema::new(window_fields))) + } else { + window_fields.extend_from_slice(&[window_expr_return_field + .as_ref() + .clone() + .with_name(fn_name)]); + Ok(Arc::new(Schema::new(window_fields))) + } +} + +/// Create a physical expression for window function +#[expect(clippy::too_many_arguments)] +pub fn create_window_expr( + fun: &WindowFunctionDefinition, + name: String, + args: &[Arc<dyn PhysicalExpr>], + partition_by: &[Arc<dyn PhysicalExpr>], + order_by: &[PhysicalSortExpr], + window_frame: Arc<WindowFrame>, + input_schema: SchemaRef, + ignore_nulls: bool, + distinct: bool, + filter: Option<Arc<dyn PhysicalExpr>>, +) -> Result<Arc<dyn WindowExpr>> { + Ok(match fun { + WindowFunctionDefinition::AggregateUDF(fun) => { + let aggregate = if distinct { + AggregateExprBuilder::new(Arc::clone(fun), args.to_vec()) + .schema(input_schema) + .alias(name) + .with_ignore_nulls(ignore_nulls) + .distinct() + .build() + .map(Arc::new)? + } else { + AggregateExprBuilder::new(Arc::clone(fun), args.to_vec()) + .schema(input_schema) + .alias(name) + .with_ignore_nulls(ignore_nulls) + .build() + .map(Arc::new)? + }; + window_expr_from_aggregate_expr( + partition_by, + order_by, + window_frame, + aggregate, + filter, + ) + } + WindowFunctionDefinition::WindowUDF(fun) => Arc::new(StandardWindowExpr::new( + create_udwf_window_expr(fun, args, &input_schema, name, ignore_nulls)?, + partition_by, + order_by, + window_frame, + )), + }) +} + +/// Creates an appropriate [`WindowExpr`] based on the window frame and +fn window_expr_from_aggregate_expr( + partition_by: &[Arc<dyn PhysicalExpr>], + order_by: &[PhysicalSortExpr], + window_frame: Arc<WindowFrame>, + aggregate: Arc<AggregateFunctionExpr>, + filter: Option<Arc<dyn PhysicalExpr>>, +) -> Arc<dyn WindowExpr> { + // Is there a potentially unlimited sized window frame? + let unbounded_window = window_frame.is_ever_expanding(); + + if !unbounded_window { + Arc::new(SlidingAggregateWindowExpr::new( + aggregate, + partition_by, + order_by, + window_frame, + filter, + )) + } else { + Arc::new(PlainAggregateWindowExpr::new( + aggregate, + partition_by, + order_by, + window_frame, + filter, + )) + } +} + +/// Creates a `StandardWindowFunctionExpr` suitable for a user defined window function +pub fn create_udwf_window_expr( + fun: &Arc<WindowUDF>, + args: &[Arc<dyn PhysicalExpr>], + input_schema: &Schema, + name: String, + ignore_nulls: bool, +) -> Result<Arc<dyn StandardWindowFunctionExpr>> { + // need to get the types into an owned vec for some reason + let input_fields: Vec<_> = args + .iter() + .map(|arg| arg.return_field(input_schema)) + .collect::<Result<_>>()?; + + let udwf_expr = Arc::new(WindowUDFExpr { + fun: Arc::clone(fun), + args: args.to_vec(), + input_fields, + name, + is_reversed: false, + ignore_nulls, + }); + + // Early validation of input expressions + // We create a partition evaluator because in the user-defined window + // implementation this is where code for parsing input expressions + // exist. The benefits are: + // - If any of the input expressions are invalid we catch them early + // in the planning phase, rather than during execution. + // - Maintains compatibility with built-in (now removed) window + // functions validation behavior. + // - Predictable and reliable error handling. + // See discussion here: + // https://github.com/apache/datafusion/pull/13201#issuecomment-2454209975 + let _ = udwf_expr.create_evaluator()?; + + Ok(udwf_expr) +} + +/// Implements [`StandardWindowFunctionExpr`] for [`WindowUDF`] +#[derive(Clone, Debug)] +pub struct WindowUDFExpr { + fun: Arc<WindowUDF>, + args: Vec<Arc<dyn PhysicalExpr>>, + /// Display name + name: String, + /// Fields of input expressions + input_fields: Vec<FieldRef>, + /// This is set to `true` only if the user-defined window function + /// expression supports evaluation in reverse order, and the + /// evaluation order is reversed. + is_reversed: bool, + /// Set to `true` if `IGNORE NULLS` is defined, `false` otherwise. + ignore_nulls: bool, +} + +impl WindowUDFExpr { + pub fn fun(&self) -> &Arc<WindowUDF> { + &self.fun + } + + /// Returns all arguments passed to this window function. + /// + /// Unlike [`StandardWindowFunctionExpr::expressions`], which returns + /// only the expressions that need batch evaluation (and may filter out + /// literal offset/default args like those for `lead`/`lag`), this + /// method returns the complete, unfiltered argument list. This is + /// needed for serialization so that all arguments survive a + /// protobuf round-trip. + pub fn args(&self) -> &[Arc<dyn PhysicalExpr>] { + &self.args + } +} + +impl StandardWindowFunctionExpr for WindowUDFExpr { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn field(&self) -> Result<FieldRef> { + self.fun + .field(WindowUDFFieldArgs::new(&self.input_fields, &self.name)) + } + + fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> { + self.fun + .expressions(ExpressionArgs::new(&self.args, &self.input_fields)) + } + + fn create_evaluator(&self) -> Result<Box<dyn PartitionEvaluator>> { + self.fun + .partition_evaluator_factory(PartitionEvaluatorArgs::new( + &self.args, + &self.input_fields, + self.is_reversed, + self.ignore_nulls, + )) + } + + fn name(&self) -> &str { + &self.name + } + + fn reverse_expr(&self) -> Option<Arc<dyn StandardWindowFunctionExpr>> { + match self.fun.reverse_expr() { + ReversedUDWF::Identical => Some(Arc::new(self.clone())), + ReversedUDWF::NotSupported => None, + ReversedUDWF::Reversed(fun) => Some(Arc::new(WindowUDFExpr { + fun, + args: self.args.clone(), + name: self.name.clone(), + input_fields: self.input_fields.clone(), + is_reversed: !self.is_reversed, + ignore_nulls: self.ignore_nulls, + })), + } + } + + fn get_result_ordering(&self, schema: &SchemaRef) -> Option<PhysicalSortExpr> { + self.fun + .sort_options() + .zip(schema.column_with_name(self.name())) + .map(|(options, (idx, field))| { + let expr = Arc::new(Column::new(field.name(), idx)); + PhysicalSortExpr { expr, options } + }) + } + + fn limit_effect(&self) -> LimitEffect { + self.fun.inner().limit_effect(self.args.as_slice()) + } +} + +pub(crate) fn calc_requirements< + T: Borrow<Arc<dyn PhysicalExpr>>, + S: Borrow<PhysicalSortExpr>, +>( + partition_by_exprs: impl IntoIterator<Item = T>, + orderby_sort_exprs: impl IntoIterator<Item = S>, +) -> Option<OrderingRequirements> { + let mut sort_reqs_with_partition = partition_by_exprs + .into_iter() + .map(|partition_by| { + PhysicalSortRequirement::new(Arc::clone(partition_by.borrow()), None) + }) + .collect::<Vec<_>>(); + let mut sort_reqs = vec![]; + for element in orderby_sort_exprs.into_iter() { + let PhysicalSortExpr { expr, options } = element.borrow(); + let sort_req = PhysicalSortRequirement::new(Arc::clone(expr), Some(*options)); + if !sort_reqs_with_partition.iter().any(|e| e.expr.eq(expr)) { + sort_reqs_with_partition.push(sort_req.clone()); + } + if !sort_reqs + .iter() + .any(|e: &PhysicalSortRequirement| e.expr.eq(expr)) + { + sort_reqs.push(sort_req); + } + } + + let mut alternatives = vec![]; + alternatives.extend(LexRequirement::new(sort_reqs_with_partition)); + alternatives.extend(LexRequirement::new(sort_reqs)); + + OrderingRequirements::new_alternatives(alternatives, false) +} + +/// This function calculates the indices such that when partition by expressions reordered with the indices +/// resulting expressions define a preset for existing ordering. +/// For instance, if input is ordered by a, b, c and PARTITION BY b, a is used, +/// this vector will be [1, 0]. It means that when we iterate b, a columns with the order [1, 0] +/// resulting vector (a, b) is a preset of the existing ordering (a, b, c). +pub fn get_ordered_partition_by_indices( + partition_by_exprs: &[Arc<dyn PhysicalExpr>], + input: &Arc<dyn ExecutionPlan>, +) -> Result<Vec<usize>> { + let (_, indices) = input + .equivalence_properties() + .find_longest_permutation(partition_by_exprs)?; + Ok(indices) +} + +pub(crate) fn get_partition_by_sort_exprs( + input: &Arc<dyn ExecutionPlan>, + partition_by_exprs: &[Arc<dyn PhysicalExpr>], + ordered_partition_by_indices: &[usize], +) -> Result<Vec<PhysicalSortExpr>> { + let ordered_partition_exprs = ordered_partition_by_indices + .iter() + .map(|idx| Arc::clone(&partition_by_exprs[*idx])) + .collect::<Vec<_>>(); + // Make sure ordered section doesn't move over the partition by expression + assert!(ordered_partition_by_indices.len() <= partition_by_exprs.len()); + let (ordering, _) = input + .equivalence_properties() + .find_longest_permutation(&ordered_partition_exprs)?; + if ordering.len() == ordered_partition_exprs.len() { + Ok(ordering) + } else { + exec_err!("Expects PARTITION BY expression to be ordered") + } +} + +pub(crate) fn window_equivalence_properties( + schema: &SchemaRef, + input: &Arc<dyn ExecutionPlan>, + window_exprs: &[Arc<dyn WindowExpr>], +) -> Result<EquivalenceProperties> { + // We need to update the schema, so we can't directly use input's equivalence + // properties. + let mut window_eq_properties = EquivalenceProperties::new(Arc::clone(schema)) + .extend(input.equivalence_properties().clone())?; + + let window_schema_len = schema.fields.len(); + let input_schema_len = window_schema_len - window_exprs.len(); + let window_expr_indices = (input_schema_len..window_schema_len).collect::<Vec<_>>(); + + for (i, expr) in window_exprs.iter().enumerate() { + let partitioning_exprs = expr.partition_by(); + let no_partitioning = partitioning_exprs.is_empty(); + + // Find "one" valid ordering for partition columns to avoid exponential complexity. + // see https://github.com/apache/datafusion/issues/17401 + let mut all_satisfied_lexs = vec![]; + let mut candidate_ordering = vec![]; + + for partition_expr in partitioning_exprs.iter() { + let sort_options = + sort_options_resolving_constant(Arc::clone(partition_expr), true); + + // Try each sort option and pick the first one that works + let mut found = false; + for sort_expr in sort_options.into_iter() { + candidate_ordering.push(sort_expr); + if let Some(lex) = LexOrdering::new(candidate_ordering.clone()) + && window_eq_properties.ordering_satisfy(lex)? + { + found = true; + break; + } + // This option didn't work, remove it and try the next one + candidate_ordering.pop(); + } + // If no sort option works for this column, we can't build a valid ordering + if !found { + candidate_ordering.clear(); + break; + } + } + + // If we successfully built an ordering for all columns, use it + // When there are no partition expressions, candidate_ordering will be empty and won't be added + if candidate_ordering.len() == partitioning_exprs.len() + && let Some(lex) = LexOrdering::new(candidate_ordering) + { + all_satisfied_lexs.push(lex); + } + // If there is a partitioning, and no possible ordering cannot satisfy + // the input plan's orderings, then we cannot further introduce any + // new orderings for the window plan. + if !no_partitioning && all_satisfied_lexs.is_empty() { + return Ok(window_eq_properties); + } else if let Some(std_expr) = expr.as_any().downcast_ref::<StandardWindowExpr>() + { + std_expr.add_equal_orderings(&mut window_eq_properties)?; + } else if let Some(plain_expr) = + expr.as_any().downcast_ref::<PlainAggregateWindowExpr>() + { + // We are dealing with plain window frames; i.e. frames having an + // unbounded starting point. + // First, check if the frame covers the whole table: + if plain_expr.get_window_frame().end_bound.is_unbounded() { + let window_col = + Arc::new(Column::new(expr.name(), i + input_schema_len)) as _; + if no_partitioning { + // Window function has a constant result across the table: + window_eq_properties + .add_constants(std::iter::once(ConstExpr::from(window_col)))? + } else { + // Window function results in a partial constant value in + // some ordering. Adjust the ordering equivalences accordingly: + let new_lexs = all_satisfied_lexs.into_iter().flat_map(|lex| { + let new_partial_consts = sort_options_resolving_constant( + Arc::clone(&window_col), + false, + ); + + new_partial_consts.into_iter().map(move |partial| { + let mut existing = lex.clone(); + existing.push(partial); + existing + }) + }); + window_eq_properties.add_orderings(new_lexs); + } + } else { + // The window frame is ever expanding, so set monotonicity comes + // into play. + plain_expr.add_equal_orderings( + &mut window_eq_properties, + window_expr_indices[i], + )?; + } + } else if let Some(sliding_expr) = + expr.as_any().downcast_ref::<SlidingAggregateWindowExpr>() + { + // We are dealing with sliding window frames; i.e. frames having an + // advancing starting point. If we have a set-monotonic expression, + // we might be able to leverage this property. + let set_monotonicity = sliding_expr.get_aggregate_expr().set_monotonicity(); + if set_monotonicity.ne(&SetMonotonicity::NotMonotonic) { + // If the window frame is ever-receding, and we have set + // monotonicity, we can utilize it to introduce new orderings. + let frame = sliding_expr.get_window_frame(); + if frame.end_bound.is_unbounded() { + let increasing = set_monotonicity.eq(&SetMonotonicity::Increasing); + let window_col = Column::new(expr.name(), i + input_schema_len); + if no_partitioning { + // Reverse set-monotonic cases with no partitioning: + window_eq_properties.add_ordering([PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(increasing, true), + )]); + } else { + // Reverse set-monotonic cases for all orderings: + for mut lex in all_satisfied_lexs.into_iter() { + lex.push(PhysicalSortExpr::new( + Arc::new(window_col.clone()), + SortOptions::new(increasing, true), + )); + window_eq_properties.add_ordering(lex); + } + } + } + // If we ensure that the elements entering the frame is greater + // than the ones leaving, and we have increasing set-monotonicity, + // then the window function result will be increasing. However, + // we also need to check if the frame is causal. If not, we cannot + // utilize set-monotonicity since the set shrinks as the frame + // boundary starts "touching" the end of the table. + else if frame.is_causal() { + // Find one valid ordering for aggregate arguments instead of + // checking all combinations + let aggregate_exprs = sliding_expr.get_aggregate_expr().expressions(); + let mut candidate_order = vec![]; + let mut asc = false; + + for (idx, expr) in aggregate_exprs.iter().enumerate() { + let mut found = false; + let sort_options = + sort_options_resolving_constant(Arc::clone(expr), false); + + // Try each option and pick the first that works + for sort_expr in sort_options.into_iter() { + let is_asc = !sort_expr.options.descending; + candidate_order.push(sort_expr); + + if let Some(lex) = LexOrdering::new(candidate_order.clone()) + && window_eq_properties.ordering_satisfy(lex)? + { + if idx == 0 { + // The first column's ordering direction determines the overall + // monotonicity behavior of the window result. + // - If the aggregate has increasing set monotonicity (e.g., MAX, COUNT) + // and the first arg is ascending, the window result is increasing + // - If the aggregate has decreasing set monotonicity (e.g., MIN) + // and the first arg is ascending, the window result is also increasing + // This flag is used to determine the final window column ordering. + asc = is_asc; + } + found = true; + break; + } + // This option didn't work, remove it and try the next one + candidate_order.pop(); + } + + // If we couldn't extend the ordering, stop trying + if !found { + break; + } + } + + // Check if we successfully built a complete ordering + let satisfied = candidate_order.len() == aggregate_exprs.len() + && !aggregate_exprs.is_empty(); + + if satisfied { + let increasing = + set_monotonicity.eq(&SetMonotonicity::Increasing); + let window_col = Column::new(expr.name(), i + input_schema_len); + if increasing && (asc || no_partitioning) { + window_eq_properties.add_ordering([PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(false, false), + )]); + } else if !increasing && (!asc || no_partitioning) { + window_eq_properties.add_ordering([PhysicalSortExpr::new( + Arc::new(window_col), + SortOptions::new(true, false), + )]); + }; + } + } + } + } + } + Ok(window_eq_properties) +} + +/// Constructs the best-fitting windowing operator (a `WindowAggExec` or a +/// `BoundedWindowExec`) for the given `input` according to the specifications +/// of `window_exprs` and `physical_partition_keys`. Here, best-fitting means +/// not requiring additional sorting and/or partitioning for the given input. +/// - A return value of `None` represents that there is no way to construct a +/// windowing operator that doesn't need additional sorting/partitioning for +/// the given input. Existing ordering should be changed to run the given +/// windowing operation. +/// - A `Some(window exec)` value contains the optimal windowing operator (a +/// `WindowAggExec` or a `BoundedWindowExec`) for the given input. +pub fn get_best_fitting_window( + window_exprs: &[Arc<dyn WindowExpr>], + input: &Arc<dyn ExecutionPlan>, + // These are the partition keys used during repartitioning. + // They are either the same with `window_expr`'s PARTITION BY columns, + // or it is empty if partitioning is not desirable for this windowing operator. + physical_partition_keys: &[Arc<dyn PhysicalExpr>], + // A [`WindowStateObserver`] installed on the source + // [`BoundedWindowAggExec`] (via [`BoundedWindowAggExec::with_state_observer`]) + // that must survive the rebuild. Ignored when the rebuilt exec is a + // [`WindowAggExec`], which does not carry an observer. `None` when the + // source is a [`WindowAggExec`] or has no observer installed. + state_observer: Option<Arc<dyn WindowStateObserver>>, +) -> Result<Option<Arc<dyn ExecutionPlan>>> { + // Contains at least one window expr and all of the partition by and order by sections + // of the window_exprs are same. + let partitionby_exprs = window_exprs[0].partition_by(); + let orderby_keys = window_exprs[0].order_by(); + let (should_reverse, input_order_mode) = + if let Some((should_reverse, input_order_mode)) = + get_window_mode(partitionby_exprs, orderby_keys, input)? + { + (should_reverse, input_order_mode) + } else { + return Ok(None); + }; + let is_unbounded = input.boundedness().is_unbounded(); + if !is_unbounded && input_order_mode != InputOrderMode::Sorted { + // Executor has bounded input and `input_order_mode` is not `InputOrderMode::Sorted` + // in this case removing the sort is not helpful, return: + return Ok(None); + }; + + let window_expr = if should_reverse { + if let Some(reversed_window_expr) = window_exprs + .iter() + .map(|e| e.get_reverse_expr()) + .collect::<Option<Vec<_>>>() + { + reversed_window_expr + } else { + // Cannot take reverse of any of the window expr + // In this case, with existing ordering window cannot be run + return Ok(None); + } + } else { + window_exprs.to_vec() + }; + + // If all window expressions can run with bounded memory, choose the + // bounded window variant: + if window_expr.iter().all(|e| e.uses_bounded_memory()) { + Ok(Some(Arc::new( + BoundedWindowAggExec::try_new( + window_expr, + Arc::clone(input), + input_order_mode, + !physical_partition_keys.is_empty(), + )? + .with_state_observer(state_observer)?, + ) as _)) + } else if input_order_mode != InputOrderMode::Sorted { + // For `WindowAggExec` to work correctly PARTITION BY columns should be sorted. + // Hence, if `input_order_mode` is not `Sorted` we should convert + // input ordering such that it can work with `Sorted` (add `SortExec`). + // Effectively `WindowAggExec` works only in `Sorted` mode. + Ok(None) + } else { + Ok(Some(Arc::new(WindowAggExec::try_new( + window_expr, + Arc::clone(input), + !physical_partition_keys.is_empty(), + )?) as _)) + } +} + +/// Compares physical ordering (output ordering of the `input` operator) with +/// `partitionby_exprs` and `orderby_keys` to decide whether existing ordering +/// is sufficient to run the current window operator. +/// - A `None` return value indicates that we can not remove the sort in question +/// (input ordering is not sufficient to run current window executor). +/// - A `Some((bool, InputOrderMode))` value indicates that the window operator +/// can run with existing input ordering, so we can remove `SortExec` before it. +/// +/// The `bool` field in the return value represents whether we should reverse window +/// operator to remove `SortExec` before it. The `InputOrderMode` field represents +/// the mode this window operator should work in to accommodate the existing ordering. +pub fn get_window_mode( + partitionby_exprs: &[Arc<dyn PhysicalExpr>], + orderby_keys: &[PhysicalSortExpr], + input: &Arc<dyn ExecutionPlan>, +) -> Result<Option<(bool, InputOrderMode)>> { + let mut input_eqs = input.equivalence_properties().clone(); + let (_, indices) = input_eqs.find_longest_permutation(partitionby_exprs)?; + let partition_by_reqs = indices + .iter() + .map(|&idx| PhysicalSortRequirement { + expr: Arc::clone(&partitionby_exprs[idx]), + options: None, + }) + .collect::<Vec<_>>(); + // Treat partition by exprs as constant. During analysis of requirements are satisfied. + let const_exprs = partitionby_exprs.iter().cloned().map(ConstExpr::from); + input_eqs.add_constants(const_exprs)?; + let reverse_orderby_keys = + orderby_keys.iter().map(|e| e.reverse()).collect::<Vec<_>>(); + for (should_swap, orderbys) in + [(false, orderby_keys), (true, reverse_orderby_keys.as_ref())] + { + let mut req = partition_by_reqs.clone(); + req.extend(orderbys.iter().cloned().map(Into::into)); + if req.is_empty() || input_eqs.ordering_satisfy_requirement(req)? { + // Window can be run with existing ordering + let mode = if indices.len() == partitionby_exprs.len() { + InputOrderMode::Sorted + } else if indices.is_empty() { + InputOrderMode::Linear + } else { + InputOrderMode::PartiallySorted(indices) + }; + return Ok(Some((should_swap, mode))); + } + } + Ok(None) +} + +/// Generates sort option variations for a given expression. +/// +/// This function is used to handle constant columns in window operations. Since constant +/// columns can be considered as having any ordering, we generate multiple sort options +/// to explore different ordering possibilities. +/// +/// # Parameters +/// - `expr`: The physical expression to generate sort options for +/// - `only_monotonic`: If false, generates all 4 possible sort options (ASC/DESC × NULLS FIRST/LAST). +/// If true, generates only 2 options that preserve set monotonicity. +/// +/// # When to use `only_monotonic = false`: +/// Use for PARTITION BY columns where we want to explore all possible orderings to find +/// one that matches the existing data ordering. +/// +/// # When to use `only_monotonic = true`: +/// Use for aggregate/window function arguments where set monotonicity needs to be preserved. +/// Only generates ASC NULLS LAST and DESC NULLS FIRST because: +/// - Set monotonicity is broken if data has increasing order but nulls come first +/// - Set monotonicity is broken if data has decreasing order but nulls come last +fn sort_options_resolving_constant( + expr: Arc<dyn PhysicalExpr>, + only_monotonic: bool, +) -> Vec<PhysicalSortExpr> { + if only_monotonic { + // Generate only the 2 options that preserve set monotonicity + vec![ + PhysicalSortExpr::new(Arc::clone(&expr), SortOptions::new(false, false)), // ASC NULLS LAST + PhysicalSortExpr::new(expr, SortOptions::new(true, true)), // DESC NULLS FIRST + ] + } else { + // Generate all 4 possible sort options for partition columns + vec![ + PhysicalSortExpr::new(Arc::clone(&expr), SortOptions::new(false, false)), // ASC NULLS LAST + PhysicalSortExpr::new(Arc::clone(&expr), SortOptions::new(false, true)), // ASC NULLS FIRST + PhysicalSortExpr::new(Arc::clone(&expr), SortOptions::new(true, false)), // DESC NULLS LAST + PhysicalSortExpr::new(expr, SortOptions::new(true, true)), // DESC NULLS FIRST + ] + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::collect; + use crate::expressions::col; + use crate::streaming::StreamingTableExec; + use crate::test::assert_is_pending; + use crate::test::exec::{BlockingExec, assert_strong_count_converges_to_zero}; + + use InputOrderMode::{Linear, PartiallySorted, Sorted}; + use arrow::compute::SortOptions; + use arrow_schema::{DataType, Field}; + use datafusion_execution::TaskContext; + use datafusion_functions_aggregate::count::count_udaf; + + use futures::FutureExt; + + fn create_test_schema() -> Result<SchemaRef> { + let nullable_column = Field::new("nullable_col", DataType::Int32, true); + let non_nullable_column = Field::new("non_nullable_col", DataType::Int32, false); + let schema = Arc::new(Schema::new(vec![nullable_column, non_nullable_column])); + + Ok(schema) + } + + fn create_test_schema2() -> Result<SchemaRef> { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, true); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, true); + let e = Field::new("e", DataType::Int32, true); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); + Ok(schema) + } + + // Generate a schema which consists of 5 columns (a, b, c, d, e) + fn create_test_schema3() -> Result<SchemaRef> { + let a = Field::new("a", DataType::Int32, true); + let b = Field::new("b", DataType::Int32, false); + let c = Field::new("c", DataType::Int32, true); + let d = Field::new("d", DataType::Int32, false); + let e = Field::new("e", DataType::Int32, false); + let schema = Arc::new(Schema::new(vec![a, b, c, d, e])); + Ok(schema) + } + + /// make PhysicalSortExpr with default options + pub fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr { + sort_expr_options(name, schema, SortOptions::default()) + } + + /// PhysicalSortExpr with specified options + pub fn sort_expr_options( + name: &str, + schema: &Schema, + options: SortOptions, + ) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col(name, schema).unwrap(), + options, + } + } + + /// Created a sorted Streaming Table exec + pub fn streaming_table_exec( + schema: &SchemaRef, + ordering: LexOrdering, + infinite_source: bool, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(StreamingTableExec::try_new( + Arc::clone(schema), + vec![], + None, + Some(ordering), + infinite_source, + None, + )?)) + } + + #[tokio::test] + async fn test_calc_requirements() -> Result<()> { + let schema = create_test_schema2()?; + let test_data = vec![ + // PARTITION BY a, ORDER BY b ASC NULLS FIRST + ( + vec!["a"], + vec![("b", true, true)], + vec![ + vec![("a", None), ("b", Some((true, true)))], + vec![("b", Some((true, true)))], + ], + ), + // PARTITION BY a, ORDER BY a ASC NULLS FIRST + ( + vec!["a"], + vec![("a", true, true)], + vec![vec![("a", None)], vec![("a", Some((true, true)))]], + ), + // PARTITION BY a, ORDER BY b ASC NULLS FIRST, c DESC NULLS LAST + ( + vec!["a"], + vec![("b", true, true), ("c", false, false)], + vec![ + vec![ + ("a", None), + ("b", Some((true, true))), + ("c", Some((false, false))), + ], + vec![("b", Some((true, true))), ("c", Some((false, false)))], + ], + ), + // PARTITION BY a, c, ORDER BY b ASC NULLS FIRST, c DESC NULLS LAST + ( + vec!["a", "c"], + vec![("b", true, true), ("c", false, false)], + vec![ + vec![("a", None), ("c", None), ("b", Some((true, true)))], + vec![("b", Some((true, true))), ("c", Some((false, false)))], + ], + ), + ]; + for (pb_params, ob_params, expected_params) in test_data { + let mut partitionbys = vec![]; + for col_name in pb_params { + partitionbys.push(col(col_name, &schema)?); + } + + let mut orderbys = vec![]; + for (col_name, descending, nulls_first) in ob_params { + let expr = col(col_name, &schema)?; + let options = SortOptions::new(descending, nulls_first); + orderbys.push(PhysicalSortExpr::new(expr, options)); + } + + let mut expected: Option<OrderingRequirements> = None; + for expected_param in expected_params.clone() { + let mut requirements = vec![]; + for (col_name, reqs) in expected_param { + let options = reqs.map(|(descending, nulls_first)| { + SortOptions::new(descending, nulls_first) + }); + let expr = col(col_name, &schema)?; + requirements.push(PhysicalSortRequirement::new(expr, options)); + } + if let Some(requirements) = LexRequirement::new(requirements) { + if let Some(alts) = expected.as_mut() { + alts.add_alternative(requirements); + } else { + expected = Some(OrderingRequirements::new(requirements)); + } + } + } + assert_eq!(calc_requirements(partitionbys, orderbys), expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, true)])); + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let window_agg_exec = Arc::new(WindowAggExec::try_new( + vec![create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count".to_owned(), + &[col("a", &schema)?], + &[], + &[], + Arc::new(WindowFrame::new(None)), + schema, + false, + false, + None, + )?], + blocking_exec, + false, + )?); + + let fut = collect(window_agg_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn get_best_fitting_window_preserves_state_observer() -> Result<()> { + // `EnforceSorting`/`EnforceDistribution` call `get_best_fitting_window` + // on a source `BoundedWindowAggExec` and replace it with the returned + // exec. Without observer propagation, a `WindowStateObserver` + // installed on the source is silently dropped by the rebuild. + use datafusion_common::ScalarValue; + use datafusion_expr::{WindowFrameBound, WindowFrameUnits}; + + struct NoopObserver; + impl WindowStateObserver for NoopObserver { + fn finalize_window_aggregate( + &self, + _partition_idx: usize, + _window_expr: &Arc<dyn WindowExpr>, + _partition_key: &datafusion_physical_expr::window::PartitionKey, + _state: Vec<ScalarValue>, + ) -> Result<()> { + Ok(()) + } + } + + let schema = create_test_schema()?; + let sort = sort_expr("nullable_col", &schema); + let ordering: LexOrdering = [sort.clone()].into(); + let source = streaming_table_exec(&schema, ordering, false)?; + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "cnt".to_string(), + &[col("nullable_col", &schema)?], + &[], + &[sort], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + source.schema(), + false, + false, + None, + )?; + + let observer: Arc<dyn WindowStateObserver> = Arc::new(NoopObserver); + let bounded = BoundedWindowAggExec::try_new( + vec![expr], + Arc::clone(&source), + Sorted, + false, + )? + .with_state_observer(Some(Arc::clone(&observer)))?; + + let rebuilt = get_best_fitting_window( + bounded.window_expr(), + bounded.input(), + &bounded.partition_keys(), + bounded.state_observer().cloned(), + )? + .expect("rebuild should produce a plan"); + let bwag = rebuilt + .downcast_ref::<BoundedWindowAggExec>() + .expect("rebuild yielded BoundedWindowAggExec"); + let installed = bwag + .state_observer() + .expect("observer preserved through rebuild"); + assert!( + Arc::ptr_eq(installed, &observer), + "observer identity preserved through rebuild", + ); + Ok(()) + } + + #[tokio::test] + async fn test_satisfy_nullable() -> Result<()> { + let schema = create_test_schema()?; + let params = vec![ + ((true, true), (false, false), false), + ((true, true), (false, true), false), + ((true, true), (true, false), false), + ((true, false), (false, true), false), + ((true, false), (false, false), false), + ((true, false), (true, true), false), + ((true, false), (true, false), true), + ]; + for ( + (physical_desc, physical_nulls_first), + (req_desc, req_nulls_first), + expected, + ) in params + { + let physical_ordering = PhysicalSortExpr { + expr: col("nullable_col", &schema)?, + options: SortOptions { + descending: physical_desc, + nulls_first: physical_nulls_first, + }, + }; + let required_ordering = PhysicalSortExpr { + expr: col("nullable_col", &schema)?, + options: SortOptions { + descending: req_desc, + nulls_first: req_nulls_first, + }, + }; + let res = physical_ordering.satisfy(&required_ordering.into(), &schema); + assert_eq!(res, expected); + } + + Ok(()) + } + + #[tokio::test] + async fn test_satisfy_non_nullable() -> Result<()> { + let schema = create_test_schema()?; + + let params = vec![ + ((true, true), (false, false), false), + ((true, true), (false, true), false), + ((true, true), (true, false), true), + ((true, false), (false, true), false), + ((true, false), (false, false), false), + ((true, false), (true, true), true), + ((true, false), (true, false), true), + ]; + for ( + (physical_desc, physical_nulls_first), + (req_desc, req_nulls_first), + expected, + ) in params + { + let physical_ordering = PhysicalSortExpr { + expr: col("non_nullable_col", &schema)?, + options: SortOptions { + descending: physical_desc, + nulls_first: physical_nulls_first, + }, + }; + let required_ordering = PhysicalSortExpr { + expr: col("non_nullable_col", &schema)?, + options: SortOptions { + descending: req_desc, + nulls_first: req_nulls_first, + }, + }; + let res = physical_ordering.satisfy(&required_ordering.into(), &schema); + assert_eq!(res, expected); + } + + Ok(()) + } + + #[tokio::test] + async fn test_get_window_mode_exhaustive() -> Result<()> { + let test_schema = create_test_schema3()?; + // Columns a,c are nullable whereas b,d are not nullable. + // Source is sorted by a ASC NULLS FIRST, b ASC NULLS FIRST, c ASC NULLS FIRST, d ASC NULLS FIRST + // Column e is not ordered. + let ordering = [ + sort_expr("a", &test_schema), + sort_expr("b", &test_schema), + sort_expr("c", &test_schema), + sort_expr("d", &test_schema), + ] + .into(); + let exec_unbounded = streaming_table_exec(&test_schema, ordering, true)?; + + // test cases consists of vector of tuples. Where each tuple represents a single test case. + // First field in the tuple is Vec<str> where each element in the vector represents PARTITION BY columns + // For instance `vec!["a", "b"]` corresponds to PARTITION BY a, b + // Second field in the tuple is Vec<str> where each element in the vector represents ORDER BY columns + // For instance, vec!["c"], corresponds to ORDER BY c ASC NULLS FIRST, (ordering is default ordering. We do not check + // for reversibility in this test). + // Third field in the tuple is Option<InputOrderMode>, which corresponds to expected algorithm mode. + // None represents that existing ordering is not sufficient to run executor with any one of the algorithms + // (We need to add SortExec to be able to run it). + // Some(InputOrderMode) represents, we can run algorithm with existing ordering; and algorithm should work in + // InputOrderMode. + let test_cases = vec![ + (vec!["a"], vec!["a"], Some(Sorted)), + (vec!["a"], vec!["b"], Some(Sorted)), + (vec!["a"], vec!["c"], None), + (vec!["a"], vec!["a", "b"], Some(Sorted)), + (vec!["a"], vec!["b", "c"], Some(Sorted)), + (vec!["a"], vec!["a", "c"], None), + (vec!["a"], vec!["a", "b", "c"], Some(Sorted)), + (vec!["b"], vec!["a"], Some(Linear)), + (vec!["b"], vec!["b"], Some(Linear)), + (vec!["b"], vec!["c"], None), + (vec!["b"], vec!["a", "b"], Some(Linear)), + (vec!["b"], vec!["b", "c"], None), + (vec!["b"], vec!["a", "c"], Some(Linear)), + (vec!["b"], vec!["a", "b", "c"], Some(Linear)), + (vec!["c"], vec!["a"], Some(Linear)), + (vec!["c"], vec!["b"], None), + (vec!["c"], vec!["c"], Some(Linear)), + (vec!["c"], vec!["a", "b"], Some(Linear)), + (vec!["c"], vec!["b", "c"], None), + (vec!["c"], vec!["a", "c"], Some(Linear)), + (vec!["c"], vec!["a", "b", "c"], Some(Linear)), + (vec!["b", "a"], vec!["a"], Some(Sorted)), + (vec!["b", "a"], vec!["b"], Some(Sorted)), + (vec!["b", "a"], vec!["c"], Some(Sorted)), + (vec!["b", "a"], vec!["a", "b"], Some(Sorted)), + (vec!["b", "a"], vec!["b", "c"], Some(Sorted)), + (vec!["b", "a"], vec!["a", "c"], Some(Sorted)), + (vec!["b", "a"], vec!["a", "b", "c"], Some(Sorted)), + (vec!["c", "b"], vec!["a"], Some(Linear)), + (vec!["c", "b"], vec!["b"], Some(Linear)), + (vec!["c", "b"], vec!["c"], Some(Linear)), + (vec!["c", "b"], vec!["a", "b"], Some(Linear)), + (vec!["c", "b"], vec!["b", "c"], Some(Linear)), + (vec!["c", "b"], vec!["a", "c"], Some(Linear)), + (vec!["c", "b"], vec!["a", "b", "c"], Some(Linear)), + (vec!["c", "a"], vec!["a"], Some(PartiallySorted(vec![1]))), + (vec!["c", "a"], vec!["b"], Some(PartiallySorted(vec![1]))), + (vec!["c", "a"], vec!["c"], Some(PartiallySorted(vec![1]))), + ( + vec!["c", "a"], + vec!["a", "b"], + Some(PartiallySorted(vec![1])), + ), + ( + vec!["c", "a"], + vec!["b", "c"], + Some(PartiallySorted(vec![1])), + ), + ( + vec!["c", "a"], + vec!["a", "c"], + Some(PartiallySorted(vec![1])), + ), + ( + vec!["c", "a"], + vec!["a", "b", "c"], + Some(PartiallySorted(vec![1])), + ), + (vec!["c", "b", "a"], vec!["a"], Some(Sorted)), + (vec!["c", "b", "a"], vec!["b"], Some(Sorted)), + (vec!["c", "b", "a"], vec!["c"], Some(Sorted)), + (vec!["c", "b", "a"], vec!["a", "b"], Some(Sorted)), + (vec!["c", "b", "a"], vec!["b", "c"], Some(Sorted)), + (vec!["c", "b", "a"], vec!["a", "c"], Some(Sorted)), + (vec!["c", "b", "a"], vec!["a", "b", "c"], Some(Sorted)), + ]; + for (case_idx, test_case) in test_cases.iter().enumerate() { + let (partition_by_columns, order_by_params, expected) = &test_case; + let mut partition_by_exprs = vec![]; + for col_name in partition_by_columns { + partition_by_exprs.push(col(col_name, &test_schema)?); + } + + let mut order_by_exprs = vec![]; + for col_name in order_by_params { + let expr = col(col_name, &test_schema)?; + // Give default ordering, this is same with input ordering direction + // In this test we do check for reversibility. + let options = SortOptions::default(); + order_by_exprs.push(PhysicalSortExpr { expr, options }); + } + let res = + get_window_mode(&partition_by_exprs, &order_by_exprs, &exec_unbounded)?; + // Since reversibility is not important in this test. Convert Option<(bool, InputOrderMode)> to Option<InputOrderMode> + let res = res.map(|(_, mode)| mode); + assert_eq!( + res, *expected, + "Unexpected result for in unbounded test case#: {case_idx:?}, case: {test_case:?}" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_get_window_mode() -> Result<()> { + let test_schema = create_test_schema3()?; + // Columns a,c are nullable whereas b,d are not nullable. + // Source is sorted by a ASC NULLS FIRST, b ASC NULLS FIRST, c ASC NULLS FIRST, d ASC NULLS FIRST + // Column e is not ordered. + let ordering = [ + sort_expr("a", &test_schema), + sort_expr("b", &test_schema), + sort_expr("c", &test_schema), + sort_expr("d", &test_schema), + ] + .into(); + let exec_unbounded = streaming_table_exec(&test_schema, ordering, true)?; + + // test cases consists of vector of tuples. Where each tuple represents a single test case. + // First field in the tuple is Vec<str> where each element in the vector represents PARTITION BY columns + // For instance `vec!["a", "b"]` corresponds to PARTITION BY a, b + // Second field in the tuple is Vec<(str, bool, bool)> where each element in the vector represents ORDER BY columns + // For instance, vec![("c", false, false)], corresponds to ORDER BY c ASC NULLS LAST, + // similarly, vec![("c", true, true)], corresponds to ORDER BY c DESC NULLS FIRST, + // Third field in the tuple is Option<(bool, InputOrderMode)>, which corresponds to expected result. + // None represents that existing ordering is not sufficient to run executor with any one of the algorithms + // (We need to add SortExec to be able to run it). + // Some((bool, InputOrderMode)) represents, we can run algorithm with existing ordering. Algorithm should work in + // InputOrderMode, bool field represents whether we should reverse window expressions to run executor with existing ordering. + // For instance, `Some((false, InputOrderMode::Sorted))`, represents that we shouldn't reverse window expressions. And algorithm + // should work in Sorted mode to work with existing ordering. + let test_cases = vec![ + // PARTITION BY a, b ORDER BY c ASC NULLS LAST + (vec!["a", "b"], vec![("c", false, false)], None), + // ORDER BY c ASC NULLS FIRST + (vec![], vec![("c", false, true)], None), + // PARTITION BY b, ORDER BY c ASC NULLS FIRST + (vec!["b"], vec![("c", false, true)], None), + // PARTITION BY a, ORDER BY c ASC NULLS FIRST + (vec!["a"], vec![("c", false, true)], None), + // PARTITION BY b, ORDER BY c ASC NULLS FIRST + ( + vec!["a", "b"], + vec![("c", false, true), ("e", false, true)], + None, + ), + // PARTITION BY a, ORDER BY b ASC NULLS FIRST + (vec!["a"], vec![("b", false, true)], Some((false, Sorted))), + // PARTITION BY a, ORDER BY a ASC NULLS FIRST + (vec!["a"], vec![("a", false, true)], Some((false, Sorted))), + // PARTITION BY a, ORDER BY a ASC NULLS LAST + (vec!["a"], vec![("a", false, false)], Some((false, Sorted))), + // PARTITION BY a, ORDER BY a DESC NULLS FIRST + (vec!["a"], vec![("a", true, true)], Some((false, Sorted))), + // PARTITION BY a, ORDER BY a DESC NULLS LAST + (vec!["a"], vec![("a", true, false)], Some((false, Sorted))), + // PARTITION BY a, ORDER BY b ASC NULLS LAST + (vec!["a"], vec![("b", false, false)], Some((false, Sorted))), + // PARTITION BY a, ORDER BY b DESC NULLS LAST + (vec!["a"], vec![("b", true, false)], Some((true, Sorted))), + // PARTITION BY a, b ORDER BY c ASC NULLS FIRST + ( + vec!["a", "b"], + vec![("c", false, true)], + Some((false, Sorted)), + ), + // PARTITION BY b, a ORDER BY c ASC NULLS FIRST + ( + vec!["b", "a"], + vec![("c", false, true)], + Some((false, Sorted)), + ), + // PARTITION BY a, b ORDER BY c DESC NULLS LAST + ( + vec!["a", "b"], + vec![("c", true, false)], + Some((true, Sorted)), + ), + // PARTITION BY e ORDER BY a ASC NULLS FIRST + ( + vec!["e"], + vec![("a", false, true)], + // For unbounded, expects to work in Linear mode. Shouldn't reverse window function. + Some((false, Linear)), + ), + // PARTITION BY b, c ORDER BY a ASC NULLS FIRST, c ASC NULLS FIRST + ( + vec!["b", "c"], + vec![("a", false, true), ("c", false, true)], + Some((false, Linear)), + ), + // PARTITION BY b ORDER BY a ASC NULLS FIRST + (vec!["b"], vec![("a", false, true)], Some((false, Linear))), + // PARTITION BY a, e ORDER BY b ASC NULLS FIRST + ( + vec!["a", "e"], + vec![("b", false, true)], + Some((false, PartiallySorted(vec![0]))), + ), + // PARTITION BY a, c ORDER BY b ASC NULLS FIRST + ( + vec!["a", "c"], + vec![("b", false, true)], + Some((false, PartiallySorted(vec![0]))), + ), + // PARTITION BY c, a ORDER BY b ASC NULLS FIRST + ( + vec!["c", "a"], + vec![("b", false, true)], + Some((false, PartiallySorted(vec![1]))), + ), + // PARTITION BY d, b, a ORDER BY c ASC NULLS FIRST + ( + vec!["d", "b", "a"], + vec![("c", false, true)], + Some((false, PartiallySorted(vec![2, 1]))), + ), + // PARTITION BY e, b, a ORDER BY c ASC NULLS FIRST + ( + vec!["e", "b", "a"], + vec![("c", false, true)], + Some((false, PartiallySorted(vec![2, 1]))), + ), + // PARTITION BY d, a ORDER BY b ASC NULLS FIRST + ( + vec!["d", "a"], + vec![("b", false, true)], + Some((false, PartiallySorted(vec![1]))), + ), + // PARTITION BY b, ORDER BY b, a ASC NULLS FIRST + ( + vec!["a"], + vec![("b", false, true), ("a", false, true)], + Some((false, Sorted)), + ), + // ORDER BY b, a ASC NULLS FIRST + (vec![], vec![("b", false, true), ("a", false, true)], None), + ]; + for (case_idx, test_case) in test_cases.iter().enumerate() { + let (partition_by_columns, order_by_params, expected) = &test_case; + let mut partition_by_exprs = vec![]; + for col_name in partition_by_columns { + partition_by_exprs.push(col(col_name, &test_schema)?); + } + + let mut order_by_exprs = vec![]; + for (col_name, descending, nulls_first) in order_by_params { + let expr = col(col_name, &test_schema)?; + let options = SortOptions { + descending: *descending, + nulls_first: *nulls_first, + }; + order_by_exprs.push(PhysicalSortExpr { expr, options }); + } + + assert_eq!( + get_window_mode(&partition_by_exprs, &order_by_exprs, &exec_unbounded)?, + *expected, + "Unexpected result for in unbounded test case#: {case_idx:?}, case: {test_case:?}" + ); + } + + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/windows/proto.rs b/native/vendor/datafusion-physical-plan/src/windows/proto.rs new file mode 100644 index 0000000..e96b0a9 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/windows/proto.rs
@@ -0,0 +1,263 @@ +// 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. + +//! Protobuf conversions shared by window execution plans. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{ + Result, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, +}; +use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, +}; +use datafusion_physical_expr::window::SlidingAggregateWindowExpr; +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; +use datafusion_proto_common::protobuf_common; +use datafusion_proto_models::protobuf::{self, physical_window_expr_node}; + +use super::{ + PlainAggregateWindowExpr, StandardWindowExpr, WindowExpr, WindowUDFExpr, + create_window_expr, schema_add_window_field, +}; + +pub(super) fn encode_physical_window_expr( + window_expr: &Arc<dyn WindowExpr>, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result<protobuf::PhysicalWindowExprNode> { + let expr = window_expr.as_any(); + let mut args = window_expr.expressions().to_vec(); + let window_frame = window_expr.get_window_frame(); + let (window_function, fun_definition, ignore_nulls, distinct) = + if let Some(plain) = expr.downcast_ref::<PlainAggregateWindowExpr>() { + let aggregate_expr = plain.get_aggregate_expr(); + ( + physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + aggregate_expr.fun().name().to_string(), + ), + ctx.encode_udaf(aggregate_expr.fun())?, + aggregate_expr.ignore_nulls(), + aggregate_expr.is_distinct(), + ) + } else if let Some(sliding) = expr.downcast_ref::<SlidingAggregateWindowExpr>() { + let aggregate_expr = sliding.get_aggregate_expr(); + ( + physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + aggregate_expr.fun().name().to_string(), + ), + ctx.encode_udaf(aggregate_expr.fun())?, + aggregate_expr.ignore_nulls(), + aggregate_expr.is_distinct(), + ) + } else if let Some(standard) = expr.downcast_ref::<StandardWindowExpr>() { + if let Some(window_udf) = standard + .get_standard_func_expr() + .as_any() + .downcast_ref::<WindowUDFExpr>() + { + // `WindowUDFExpr::args` returns the full, unfiltered argument list so + // every argument survives the round-trip. + args = window_udf.args().to_vec(); + ( + physical_window_expr_node::WindowFunction::UserDefinedWindowFunction( + window_udf.fun().name().to_string(), + ), + ctx.encode_udwf(window_udf.fun().as_ref())?, + false, + false, + ) + } else { + return not_impl_err!( + "User-defined window function not supported: {window_expr:?}" + ); + } + } else { + return not_impl_err!("WindowExpr not supported: {window_expr:?}"); + }; + + let args = ctx.encode_expressions(&args)?; + let partition_by = ctx.encode_expressions(window_expr.partition_by())?; + let order_by = sort_exprs_try_to_proto(window_expr.order_by(), &ctx.expr_ctx())?; + + Ok(protobuf::PhysicalWindowExprNode { + args, + partition_by, + order_by, + window_frame: Some(encode_window_frame(window_frame.as_ref())?), + window_function: Some(window_function), + name: window_expr.name().to_string(), + fun_definition, + ignore_nulls, + distinct, + }) +} + +pub(super) fn decode_physical_window_expr( + proto: &protobuf::PhysicalWindowExprNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + input_schema: &Schema, +) -> Result<Arc<dyn WindowExpr>> { + let args = proto + .args + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema)) + .collect::<Result<Vec<_>>>()?; + let partition_by = proto + .partition_by + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema)) + .collect::<Result<Vec<_>>>()?; + let order_by = + sort_exprs_try_from_proto(&proto.order_by, &ctx.expr_ctx(input_schema))?; + let window_frame = proto + .window_frame + .as_ref() + .map(decode_window_frame) + .transpose()? + .ok_or_else(|| { + internal_datafusion_err!("Missing required field 'window_frame' in protobuf") + })?; + let function = match proto.window_function.as_ref() { + Some(physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + name, + )) => WindowFunctionDefinition::AggregateUDF( + ctx.decode_udaf(name, proto.fun_definition.as_deref())?, + ), + Some(physical_window_expr_node::WindowFunction::UserDefinedWindowFunction( + name, + )) => WindowFunctionDefinition::WindowUDF( + ctx.decode_udwf(name, proto.fun_definition.as_deref())?, + ), + None => { + return internal_err!("Missing required field 'window_function' in protobuf"); + } + }; + + let name = proto.name.clone(); + // TODO: Remove extended_schema if functions are all UDAF + let extended_schema = schema_add_window_field(&args, input_schema, &function, &name)?; + create_window_expr( + &function, + name, + &args, + &partition_by, + &order_by, + Arc::new(window_frame), + extended_schema, + proto.ignore_nulls, + proto.distinct, + None, + ) +} + +fn encode_window_frame(window_frame: &WindowFrame) -> Result<protobuf::WindowFrame> { + let units = match window_frame.units { + WindowFrameUnits::Rows => protobuf::WindowFrameUnits::Rows, + WindowFrameUnits::Range => protobuf::WindowFrameUnits::Range, + WindowFrameUnits::Groups => protobuf::WindowFrameUnits::Groups, + }; + Ok(protobuf::WindowFrame { + window_frame_units: units.into(), + start_bound: Some(encode_window_frame_bound(&window_frame.start_bound)?), + end_bound: Some(protobuf::window_frame::EndBound::Bound( + encode_window_frame_bound(&window_frame.end_bound)?, + )), + }) +} + +fn encode_window_frame_bound( + bound: &WindowFrameBound, +) -> Result<protobuf::WindowFrameBound> { + let encode_value = |value: &ScalarValue| -> Result<protobuf_common::ScalarValue> { + Ok(value.try_into()?) + }; + Ok(match bound { + WindowFrameBound::CurrentRow => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow.into(), + bound_value: None, + }, + WindowFrameBound::Preceding(value) => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::Preceding.into(), + bound_value: Some(encode_value(value)?), + }, + WindowFrameBound::Following(value) => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::Following.into(), + bound_value: Some(encode_value(value)?), + }, + }) +} + +fn decode_window_frame(window_frame: &protobuf::WindowFrame) -> Result<WindowFrame> { + let units = protobuf::WindowFrameUnits::try_from(window_frame.window_frame_units) + .map_err(|_| { + internal_datafusion_err!( + "Received a WindowFrame message with unknown WindowFrameUnits {}", + window_frame.window_frame_units + ) + })?; + let units = match units { + protobuf::WindowFrameUnits::Rows => WindowFrameUnits::Rows, + protobuf::WindowFrameUnits::Range => WindowFrameUnits::Range, + protobuf::WindowFrameUnits::Groups => WindowFrameUnits::Groups, + }; + let start_bound = + decode_window_frame_bound(window_frame.start_bound.as_ref().ok_or_else( + || internal_datafusion_err!("Missing start_bound in WindowFrame"), + )?)?; + let end_bound = window_frame + .end_bound + .as_ref() + .map(|end_bound| match end_bound { + protobuf::window_frame::EndBound::Bound(bound) => { + decode_window_frame_bound(bound) + } + }) + .transpose()? + .unwrap_or(WindowFrameBound::CurrentRow); + Ok(WindowFrame::new_bounds(units, start_bound, end_bound)) +} + +fn decode_window_frame_bound( + bound: &protobuf::WindowFrameBound, +) -> Result<WindowFrameBound> { + let decode_value = |value: &protobuf_common::ScalarValue| -> Result<ScalarValue> { + Ok(ScalarValue::try_from(value)?) + }; + let bound_type = protobuf::WindowFrameBoundType::try_from( + bound.window_frame_bound_type, + ) + .map_err(|_| { + internal_datafusion_err!( + "Received a WindowFrameBound message with unknown WindowFrameBoundType {}", + bound.window_frame_bound_type + ) + })?; + match bound_type { + protobuf::WindowFrameBoundType::CurrentRow => Ok(WindowFrameBound::CurrentRow), + protobuf::WindowFrameBoundType::Preceding => match &bound.bound_value { + Some(value) => Ok(WindowFrameBound::Preceding(decode_value(value)?)), + None => Ok(WindowFrameBound::Preceding(ScalarValue::UInt64(None))), + }, + protobuf::WindowFrameBoundType::Following => match &bound.bound_value { + Some(value) => Ok(WindowFrameBound::Following(decode_value(value)?)), + None => Ok(WindowFrameBound::Following(ScalarValue::UInt64(None))), + }, + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/windows/utils.rs b/native/vendor/datafusion-physical-plan/src/windows/utils.rs new file mode 100644 index 0000000..be38976 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/windows/utils.rs
@@ -0,0 +1,37 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::datatypes::{Schema, SchemaBuilder}; +use datafusion_common::Result; +use datafusion_physical_expr::window::WindowExpr; +use std::sync::Arc; + +pub(crate) fn create_schema( + input_schema: &Schema, + window_expr: &[Arc<dyn WindowExpr>], +) -> Result<Schema> { + let capacity = input_schema.fields().len() + window_expr.len(); + let mut builder = SchemaBuilder::with_capacity(capacity); + builder.extend(input_schema.fields().iter().cloned()); + // append results to the schema + for expr in window_expr { + builder.push(expr.field()?); + } + Ok(builder + .finish() + .with_metadata(input_schema.metadata().clone())) +}
diff --git a/native/vendor/datafusion-physical-plan/src/windows/window_agg_exec.rs b/native/vendor/datafusion-physical-plan/src/windows/window_agg_exec.rs new file mode 100644 index 0000000..d794e7d --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/windows/window_agg_exec.rs
@@ -0,0 +1,678 @@ +// 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. + +//! Stream and channel implementations for window function expressions. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +#[cfg(feature = "proto")] +use super::proto::{decode_physical_window_expr, encode_physical_window_expr}; +use super::utils::create_schema; +use crate::execution_plan::{CardinalityEffect, EmissionType}; +use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::stream::EmptyRecordBatchStream; +use crate::windows::{ + calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, + window_equivalence_properties, +}; +use crate::{ + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + Statistics, WindowExpr, validate_child_count, +}; + +use arrow::array::ArrayRef; +use arrow::compute::{concat, concat_batches}; +use arrow::datatypes::SchemaRef; +use arrow::error::ArrowError; +use arrow::record_batch::RecordBatch; +use datafusion_common::stats::Precision; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::utils::{evaluate_partition_ranges, transpose}; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr_common::sort_expr::{ + OrderingRequirements, PhysicalSortExpr, +}; + +use futures::{Stream, StreamExt, ready}; + +/// Window execution plan +#[derive(Debug, Clone)] +pub struct WindowAggExec { + /// Input plan + pub(crate) input: Arc<dyn ExecutionPlan>, + /// Window function expression + window_expr: Vec<Arc<dyn WindowExpr>>, + /// Schema after the window is run + schema: SchemaRef, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Partition by indices that defines preset for existing ordering + // see `get_ordered_partition_by_indices` for more details. + ordered_partition_by_indices: Vec<usize>, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, + /// If `can_partition` is false, partition_keys is always empty. + can_repartition: bool, +} + +impl WindowAggExec { + /// Create a new execution plan for window aggregates + pub fn try_new( + window_expr: Vec<Arc<dyn WindowExpr>>, + input: Arc<dyn ExecutionPlan>, + can_repartition: bool, + ) -> Result<Self> { + let schema = create_schema(&input.schema(), &window_expr)?; + let schema = Arc::new(schema); + + let ordered_partition_by_indices = + get_ordered_partition_by_indices(window_expr[0].partition_by(), &input)?; + let cache = Self::compute_properties(&schema, &input, &window_expr)?; + Ok(Self { + input, + window_expr, + schema, + metrics: ExecutionPlanMetricsSet::new(), + ordered_partition_by_indices, + cache: Arc::new(cache), + can_repartition, + }) + } + + /// Window expressions + pub fn window_expr(&self) -> &[Arc<dyn WindowExpr>] { + &self.window_expr + } + + /// Input plan + pub fn input(&self) -> &Arc<dyn ExecutionPlan> { + &self.input + } + + /// Return the output sort order of partition keys: For example + /// OVER(PARTITION BY a, ORDER BY b) -> would give sorting of the column a + // We are sure that partition by columns are always at the beginning of sort_keys + // Hence returned `PhysicalSortExpr` corresponding to `PARTITION BY` columns can be used safely + // to calculate partition separation points + pub fn partition_by_sort_keys(&self) -> Result<Vec<PhysicalSortExpr>> { + let partition_by = self.window_expr()[0].partition_by(); + get_partition_by_sort_exprs( + &self.input, + partition_by, + &self.ordered_partition_by_indices, + ) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: &SchemaRef, + input: &Arc<dyn ExecutionPlan>, + window_exprs: &[Arc<dyn WindowExpr>], + ) -> Result<PlanProperties> { + // Calculate equivalence properties: + let eq_properties = window_equivalence_properties(schema, input, window_exprs)?; + + // Get output partitioning: + // Because we can have repartitioning using the partition keys this + // would be either 1 or more than 1 depending on the presence of repartitioning. + let output_partitioning = input.output_partitioning().clone(); + + // Construct properties cache: + Ok(PlanProperties::new( + eq_properties, + output_partitioning, + // TODO: Emission type and boundedness information can be enhanced here + EmissionType::Final, + input.boundedness(), + )) + } + + pub fn partition_keys(&self) -> Vec<Arc<dyn PhysicalExpr>> { + if !self.can_repartition { + vec![] + } else { + let all_partition_keys = self + .window_expr() + .iter() + .map(|expr| expr.partition_by().to_vec()) + .collect::<Vec<_>>(); + + all_partition_keys + .into_iter() + .min_by_key(|s| s.len()) + .unwrap_or_else(Vec::new) + } + } +} + +impl DisplayAs for WindowAggExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "WindowAggExec: ")?; + let g: Vec<String> = self + .window_expr + .iter() + .map(|e| { + format!( + "{}: {:?}, frame: {:?}", + e.name().to_owned(), + e.field(), + e.get_window_frame() + ) + }) + .collect(); + write!(f, "wdw=[{}]", g.join(", "))?; + } + DisplayFormatType::TreeRender => { + let g: Vec<String> = self + .window_expr + .iter() + .map(|e| e.name().to_owned().to_string()) + .collect(); + writeln!(f, "select_list={}", g.join(", "))?; + } + } + Ok(()) + } +} + +impl ExecutionPlan for WindowAggExec { + fn name(&self) -> &'static str { + "WindowAggExec" + } + + /// Return a reference to Any that can be used for downcasting + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + let expressions = self.window_expr.iter().flat_map(|window_expr| { + let expressions = window_expr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.partition_by_exprs) + .chain(expressions.order_by_exprs) + }); + crate::apply_expression_roots(expressions, f) + } + + fn maintains_input_order(&self) -> Vec<bool> { + vec![true] + } + + fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> { + let partition_bys = self.window_expr()[0].partition_by(); + let order_keys = self.window_expr()[0].order_by(); + if self.ordered_partition_by_indices.len() < partition_bys.len() { + vec![calc_requirements(partition_bys, order_keys)] + } else { + let partition_bys = self + .ordered_partition_by_indices + .iter() + .map(|idx| &partition_bys[*idx]); + vec![calc_requirements(partition_bys, order_keys)] + } + } + + fn required_input_distribution(&self) -> Vec<Distribution> { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + if self.partition_keys().is_empty() { + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) + } else { + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) + } + } + + fn replace_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + options: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(WindowAggExec::try_new( + self.window_expr.clone(), + children.swap_remove(0), + true, + )?)), + } + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let input = self.input.execute(partition, context)?; + let stream = Box::pin(WindowAggStream::new( + Arc::clone(&self.schema), + self.window_expr.clone(), + input, + BaselineMetrics::new(&self.metrics, partition), + self.partition_by_sort_keys()?, + self.ordered_partition_by_indices.clone(), + )?); + Ok(stream) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn child_stats_requests(&self, partition: Option<usize>) -> Vec<ChildStats> { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + let input_stat = input_stats[0].as_ref().clone(); + let win_cols = self.window_expr.len(); + let input_cols = self.input.schema().fields().len(); + // TODO stats: some windowing function will maintain invariants such as min, max... + let mut column_statistics = Vec::with_capacity(win_cols + input_cols); + // copy stats of the input to the beginning of the schema. + column_statistics.extend(input_stat.column_statistics); + for _ in 0..win_cols { + column_statistics.push(ColumnStatistics::new_unknown()) + } + Ok(Arc::new(Statistics { + num_rows: input_stat.num_rows, + column_statistics, + total_byte_size: Precision::Absent, + })) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `WindowAggExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + input, + window_expr, + // Derived at construction by `create_schema` from the input schema + // and the window expressions. + schema: _, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + // Derived at construction by `get_ordered_partition_by_indices`. + ordered_partition_by_indices: _, + // Derived at construction by `Self::compute_properties`. + cache: _, + // No wire field of its own; it is folded into `partition_keys` + // below, since `partition_keys()` returns an empty vec when this is + // false and the decoder recovers it as `!partition_keys.is_empty()`. + can_repartition: _, + } = self; + + let input = ctx.encode_child(input)?; + let window_expr = window_expr + .iter() + .map(|expr| encode_physical_window_expr(expr, ctx)) + .collect::<Result<Vec<_>>>()?; + let partition_keys = self + .partition_keys() + .iter() + .map(|expr| ctx.encode_expr(expr)) + .collect::<Result<Vec<_>>>()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new( + protobuf::WindowAggExecNode { + input: Some(Box::new(input)), + window_expr, + partition_keys, + // `None` distinguishes a `WindowAggExec` from a + // `BoundedWindowAggExec` on the shared `Window` variant. + input_order_mode: None, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl WindowAggExec { + /// Reconstruct a window plan from its protobuf representation. + /// + /// This returns a [`WindowAggExec`] when `input_order_mode` is absent and a + /// [`BoundedWindowAggExec`] when it is present. + /// + /// [`BoundedWindowAggExec`]: crate::windows::BoundedWindowAggExec + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result<Arc<dyn ExecutionPlan>> { + use super::BoundedWindowAggExec; + use crate::InputOrderMode; + use datafusion_proto_models::protobuf; + use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode; + + let window_agg = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Window, + "WindowAggExec", + ); + // Exhaustive destructure: a new field on `WindowAggExecNode` is a + // compile error here rather than a silently ignored wire field. + let protobuf::WindowAggExecNode { + input, + window_expr, + partition_keys, + input_order_mode, + } = window_agg.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "WindowAggExec", "input")?; + let input_schema = input.schema(); + let window_expr = window_expr + .iter() + .map(|expr| decode_physical_window_expr(expr, ctx, input_schema.as_ref())) + .collect::<Result<Vec<_>>>()?; + let partition_keys = partition_keys + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) + .collect::<Result<Vec<_>>>()?; + + if let Some(input_order_mode) = input_order_mode.as_ref() { + let input_order_mode = match input_order_mode { + ProtoInputOrderMode::Linear(_) => InputOrderMode::Linear, + ProtoInputOrderMode::PartiallySorted( + protobuf::PartiallySortedInputOrderMode { columns }, + ) => InputOrderMode::PartiallySorted( + columns.iter().map(|column| *column as usize).collect(), + ), + ProtoInputOrderMode::Sorted(_) => InputOrderMode::Sorted, + }; + Ok(Arc::new(BoundedWindowAggExec::try_new( + window_expr, + input, + input_order_mode, + // `can_repartition` has no wire field: the encoder writes an + // empty `partition_keys` when it is false. + !partition_keys.is_empty(), + )?)) + } else { + Ok(Arc::new(WindowAggExec::try_new( + window_expr, + input, + // See above: `can_repartition` is recovered from `partition_keys`. + !partition_keys.is_empty(), + )?)) + } + } +} + +/// Compute the window aggregate columns +fn compute_window_aggregates( + window_expr: &[Arc<dyn WindowExpr>], + batch: &RecordBatch, +) -> Result<Vec<ArrayRef>> { + window_expr + .iter() + .map(|window_expr| window_expr.evaluate(batch)) + .collect() +} + +/// stream for window aggregation plan +pub struct WindowAggStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + batches: Vec<RecordBatch>, + finished: bool, + window_expr: Vec<Arc<dyn WindowExpr>>, + partition_by_sort_keys: Vec<PhysicalSortExpr>, + baseline_metrics: BaselineMetrics, + ordered_partition_by_indices: Vec<usize>, +} + +impl WindowAggStream { + /// Create a new WindowAggStream + pub fn new( + schema: SchemaRef, + window_expr: Vec<Arc<dyn WindowExpr>>, + input: SendableRecordBatchStream, + baseline_metrics: BaselineMetrics, + partition_by_sort_keys: Vec<PhysicalSortExpr>, + ordered_partition_by_indices: Vec<usize>, + ) -> Result<Self> { + // In WindowAggExec all partition by columns should be ordered. + assert_eq_or_internal_err!( + window_expr[0].partition_by().len(), + ordered_partition_by_indices.len(), + "All partition by columns should have an ordering" + ); + Ok(Self { + schema, + input, + batches: vec![], + finished: false, + window_expr, + baseline_metrics, + partition_by_sort_keys, + ordered_partition_by_indices, + }) + } + + fn compute_aggregates(&self) -> Result<Option<RecordBatch>> { + // record compute time on drop + let _timer = self.baseline_metrics.elapsed_compute().timer(); + + let batch = concat_batches(&self.input.schema(), &self.batches)?; + if batch.num_rows() == 0 { + return Ok(None); + } + + let partition_by_sort_keys = self + .ordered_partition_by_indices + .iter() + .map(|idx| self.partition_by_sort_keys[*idx].evaluate_to_sort_column(&batch)) + .collect::<Result<Vec<_>>>()?; + let partition_points = + evaluate_partition_ranges(batch.num_rows(), &partition_by_sort_keys)?; + + let mut partition_results = vec![]; + // Calculate window cols + for partition_point in partition_points { + let length = partition_point.end - partition_point.start; + partition_results.push(compute_window_aggregates( + &self.window_expr, + &batch.slice(partition_point.start, length), + )?) + } + let columns = transpose(partition_results) + .iter() + .map(|elems| concat(&elems.iter().map(|x| x.as_ref()).collect::<Vec<_>>())) + .collect::<Vec<_>>() + .into_iter() + .collect::<Result<Vec<ArrayRef>, ArrowError>>()?; + + // combine with the original cols + // note the setup of window aggregates is that they newly calculated window + // expression results are always appended to the columns + let mut batch_columns = batch.columns().to_vec(); + // calculate window cols + batch_columns.extend_from_slice(&columns); + Ok(Some(RecordBatch::try_new( + Arc::clone(&self.schema), + batch_columns, + )?)) + } +} + +impl Stream for WindowAggStream { + type Item = Result<RecordBatch>; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll<Option<Self::Item>> { + let poll = self.poll_next_inner(cx); + self.baseline_metrics.record_poll(poll) + } +} + +impl WindowAggStream { + #[inline] + fn poll_next_inner( + &mut self, + cx: &mut Context<'_>, + ) -> Poll<Option<Result<RecordBatch>>> { + if self.finished { + return Poll::Ready(None); + } + + loop { + return Poll::Ready(Some(match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + self.batches.push(batch); + continue; + } + Some(Err(e)) => Err(e), + None => { + // Release the input pipeline's resources before computing + // the final aggregates. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + let Some(result) = self.compute_aggregates()? else { + return Poll::Ready(None); + }; + self.finished = true; + // Empty record batches should not be emitted. + // They need to be treated as [`Option<RecordBatch>`]es and handled separately + debug_assert!(result.num_rows() > 0); + Ok(result) + } + })); + } + } +} + +impl RecordBatchStream for WindowAggStream { + /// Get the schema + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test::TestMemoryExec; + use crate::windows::create_window_expr; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::ScalarValue; + use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, + }; + use datafusion_functions_aggregate::count::count_udaf; + + #[test] + fn test_window_agg_cardinality_effect() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])); + let input: Arc<dyn ExecutionPlan> = + Arc::new(TestMemoryExec::try_new(&[], Arc::clone(&schema), None)?); + let args = vec![crate::expressions::col("a", &schema)?]; + let window_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count(a)".to_string(), + &args, + &[], + &[], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + + let window = WindowAggExec::try_new(vec![window_expr], input, true)?; + assert!(matches!( + window.cardinality_effect(), + CardinalityEffect::Equal + )); + Ok(()) + } +}
diff --git a/native/vendor/datafusion-physical-plan/src/work_table.rs b/native/vendor/datafusion-physical-plan/src/work_table.rs new file mode 100644 index 0000000..b5d6fd4 --- /dev/null +++ b/native/vendor/datafusion-physical-plan/src/work_table.rs
@@ -0,0 +1,375 @@ +// 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. + +//! Defines the work table query plan + +use std::any::Any; +use std::sync::{Arc, Mutex}; + +use crate::coop::cooperative; +use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; +use crate::memory::MemoryStream; +use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, +}; + +use crate::statistics::StatisticsArgs; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; + +/// A vector of record batches with a memory reservation. +#[derive(Debug)] +pub(super) struct ReservedBatches { + batches: Vec<RecordBatch>, + reservation: MemoryReservation, +} + +impl ReservedBatches { + pub(super) fn new(batches: Vec<RecordBatch>, reservation: MemoryReservation) -> Self { + ReservedBatches { + batches, + reservation, + } + } +} + +/// The name is from PostgreSQL's terminology. +/// See <https://wiki.postgresql.org/wiki/CTEReadme#How_Recursion_Works> +/// This table serves as a mirror or buffer between each iteration of a recursive query. +#[derive(Debug)] +pub struct WorkTable { + batches: Mutex<Option<ReservedBatches>>, + name: String, +} + +impl WorkTable { + /// Create a new work table. + pub(super) fn new(name: String) -> Self { + Self { + batches: Mutex::new(None), + name, + } + } + + /// Take the previously written batches from the work table. + /// This will be called by the [`WorkTableExec`] when it is executed. + fn take(&self) -> Result<ReservedBatches> { + self.batches + .lock() + .unwrap() + .take() + .ok_or_else(|| internal_datafusion_err!("Unexpected empty work table")) + } + + /// Update the results of a recursive query iteration to the work table. + pub(super) fn update(&self, batches: ReservedBatches) { + self.batches.lock().unwrap().replace(batches); + } +} + +/// A temporary "working table" operation where the input data will be +/// taken from the named handle during the execution and will be re-published +/// as is (kind of like a mirror). +/// +/// Most notably used in the implementation of recursive queries where the +/// underlying relation does not exist yet but the data will come as the previous +/// term is evaluated. This table will be used such that the recursive plan +/// will register a receiver in the task context and this plan will use that +/// receiver to get the data and stream it back up so that the batches are available +/// in the next iteration. +#[derive(Clone, Debug)] +pub struct WorkTableExec { + /// Name of the relation handler + name: String, + /// The schema of the stream + schema: SchemaRef, + /// Projection to apply to build the output stream from the recursion state + projection: Option<Vec<usize>>, + /// The work table + work_table: Arc<WorkTable>, + /// Execution metrics + metrics: ExecutionPlanMetricsSet, + /// Cache holding plan properties like equivalences, output partitioning etc. + cache: Arc<PlanProperties>, +} + +impl WorkTableExec { + /// Create a new execution plan for a worktable exec. + pub fn new( + name: String, + mut schema: SchemaRef, + projection: Option<Vec<usize>>, + ) -> Result<Self> { + if let Some(projection) = &projection { + schema = Arc::new(schema.project(projection)?); + } + let cache = Self::compute_properties(Arc::clone(&schema)); + Ok(Self { + name: name.clone(), + schema, + projection, + work_table: Arc::new(WorkTable::new(name)), + metrics: ExecutionPlanMetricsSet::new(), + cache: Arc::new(cache), + }) + } + + /// Ref to name + pub fn name(&self) -> &str { + &self.name + } + + /// Arc clone of ref to schema + pub fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ) + .with_scheduling_type(SchedulingType::Cooperative) + } +} + +impl DisplayAs for WorkTableExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "WorkTableExec: name={}", self.name) + } + DisplayFormatType::TreeRender => { + write!(f, "name={}", self.name) + } + } + } +} + +impl ExecutionPlan for WorkTableExec { + fn name(&self) -> &'static str { + "WorkTableExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.cache + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc<Self>, + _: Vec<Arc<dyn ExecutionPlan>>, + _: ReplaceChildrenOptions, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::clone(&self) as Arc<dyn ExecutionPlan>) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + /// Stream the batches that were written to the work table. + fn execute( + &self, + partition: usize, + _context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + // WorkTable streams must be the plan base. + assert_eq_or_internal_err!( + partition, + 0, + "WorkTableExec got an invalid partition {partition} (expected 0)" + ); + let ReservedBatches { + mut batches, + reservation, + } = self.work_table.take()?; + if let Some(projection) = &self.projection { + // We apply the projection + // TODO: it would be better to apply it as soon as possible and not only here + // TODO: an aggressive projection makes the memory reservation smaller, even if we do not edit it + batches = batches + .into_iter() + .map(|b| b.project(projection)) + .collect::<Result<Vec<_>, _>>()?; + } + + let stream = MemoryStream::try_new(batches, Arc::clone(&self.schema), None)? + .with_reservation(reservation); + Ok(Box::pin(cooperative(stream))) + } + + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc<Statistics>], + _args: &StatisticsArgs, + ) -> Result<Arc<Statistics>> { + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } + + /// Injects run-time state into this `WorkTableExec`. + /// + /// The only state this node currently understands is an [`Arc<WorkTable>`]. + /// If `state` can be down-cast to that type, a new `WorkTableExec` backed + /// by the provided work table is returned. Otherwise `None` is returned + /// so that callers can attempt to propagate the state further down the + /// execution plan tree. + fn with_new_state( + &self, + state: Arc<dyn Any + Send + Sync>, + ) -> Option<Arc<dyn ExecutionPlan>> { + // Down-cast to the expected state type; propagate `None` on failure + let work_table = state.downcast::<WorkTable>().ok()?; + + if work_table.name != self.name { + return None; // Different table + } + + Some(Arc::new(Self { + name: self.name.clone(), + schema: Arc::clone(&self.schema), + projection: self.projection.clone(), + metrics: ExecutionPlanMetricsSet::new(), + work_table, + cache: Arc::clone(&self.cache), + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, Int16Array, Int32Array, Int64Array}; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_execution::memory_pool::{MemoryConsumer, UnboundedMemoryPool}; + use futures::StreamExt; + + #[test] + fn test_work_table() { + let work_table = WorkTable::new("test".into()); + // Can't take from empty work_table + assert!(work_table.take().is_err()); + + let pool = Arc::new(UnboundedMemoryPool::default()) as _; + let reservation = MemoryConsumer::new("test_work_table").register(&pool); + + // Update batch to work_table + let array: ArrayRef = Arc::new((0..5).collect::<Int32Array>()); + let batch = RecordBatch::try_from_iter(vec![("col", array)]).unwrap(); + reservation.try_grow(100).unwrap(); + work_table.update(ReservedBatches::new(vec![batch.clone()], reservation)); + // Take from work_table + let reserved_batches = work_table.take().unwrap(); + assert_eq!(reserved_batches.batches, vec![batch.clone()]); + + // Consume the batch by the MemoryStream + let memory_stream = + MemoryStream::try_new(reserved_batches.batches, batch.schema(), None) + .unwrap() + .with_reservation(reserved_batches.reservation); + + // Should still be reserved + assert_eq!(pool.reserved(), 100); + + // The reservation should be freed after drop the memory_stream + drop(memory_stream); + assert_eq!(pool.reserved(), 0); + } + + #[tokio::test] + async fn test_work_table_exec() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int16, false), + ])); + let work_table_exec = + WorkTableExec::new("wt".into(), Arc::clone(&schema), Some(vec![2, 1])) + .unwrap(); + + // We inject the work table + let work_table = Arc::new(WorkTable::new("wt".into())); + let work_table_exec = work_table_exec + .with_new_state(Arc::clone(&work_table) as _) + .unwrap(); + + // We update the work table + let pool = Arc::new(UnboundedMemoryPool::default()) as _; + let reservation = MemoryConsumer::new("test_work_table").register(&pool); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(Int16Array::from(vec![1, 2, 3, 4, 5])), + ], + ) + .unwrap(); + work_table.update(ReservedBatches::new(vec![batch], reservation)); + + // We get back the batch from the work table + let returned_batch = work_table_exec + .execute(0, Arc::new(TaskContext::default())) + .unwrap() + .next() + .await + .unwrap() + .unwrap(); + assert_eq!( + returned_batch, + RecordBatch::try_from_iter(vec![ + ("c", Arc::new(Int16Array::from(vec![1, 2, 3, 4, 5])) as _), + ("b", Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _), + ]) + .unwrap() + ); + } +}
diff --git a/native/vendor/datafusion-pruning/Cargo.toml b/native/vendor/datafusion-pruning/Cargo.toml new file mode 100644 index 0000000..7d46c3d --- /dev/null +++ b/native/vendor/datafusion-pruning/Cargo.toml
@@ -0,0 +1,126 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2024" +name = "datafusion-pruning" +version = "55.1.0" +authors = ["Apache DataFusion <dev@datafusion.apache.org>"] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "DataFusion Pruning Logic" +homepage = "https://datafusion.apache.org" +readme = "README.md" +license = "Apache-2.0" +repository = "https://github.com/apache/datafusion" +resolver = "2" + +[lib] +name = "datafusion_pruning" +path = "src/lib.rs" + +[dependencies.arrow] +version = "59.2.0" +features = [ + "prettyprint", + "chrono-tz", +] + +[dependencies.datafusion-common] +version = "55.1.0" +default-features = true + +[dependencies.datafusion-datasource] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-expr-common] +version = "55.1.0" + +[dependencies.datafusion-physical-expr] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-physical-expr-common] +version = "55.1.0" +default-features = false + +[dependencies.datafusion-physical-plan] +version = "55.1.0" + +[dependencies.log] +version = "^0.4" + +[dev-dependencies.datafusion-expr] +version = "55.1.0" +default-features = false + +[dev-dependencies.datafusion-functions-nested] +version = "55.1.0" +default-features = false + +[dev-dependencies.insta] +version = "1.47.2" +features = [ + "glob", + "filters", +] + +[dev-dependencies.itertools] +version = "0.15" + +[dev-dependencies.rstest] +version = "0.26.1" + +[lints.clippy] +allow_attributes = "warn" +assigning_clones = "warn" +inefficient_to_string = "warn" +large_futures = "warn" +needless_pass_by_value = "warn" +or_fun_call = "warn" +uninlined_format_args = "warn" +unnecessary_lazy_evaluations = "warn" +unused_async = "warn" +used_underscore_binding = "warn" + +[lints.rust] +unused_qualifications = "deny" + +[lints.rust.unexpected_cfgs] +level = "warn" +priority = 0 +check-cfg = [ + 'cfg(datafusion_coop, values("tokio", "tokio_fallback", "per_stream"))', + "cfg(coverage)", + "cfg(coverage_nightly)", +]
diff --git a/native/vendor/datafusion-pruning/Cargo.toml.orig b/native/vendor/datafusion-pruning/Cargo.toml.orig new file mode 100644 index 0000000..7558640 --- /dev/null +++ b/native/vendor/datafusion-pruning/Cargo.toml.orig
@@ -0,0 +1,49 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-pruning" +description = "DataFusion Pruning Logic" +readme = "README.md" +version = { workspace = true } +edition = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +authors = { workspace = true } + +# Note: add additional linter rules in lib.rs. +# Rust does not support workspace + new linter rules in subcrates yet +# https://github.com/rust-lang/cargo/issues/13157 +[lints] +workspace = true + +[dependencies] +arrow = { workspace = true } +datafusion-common = { workspace = true, default-features = true } +datafusion-datasource = { workspace = true } +datafusion-expr-common = { workspace = true, default-features = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-expr-common = { workspace = true } +datafusion-physical-plan = { workspace = true } +log = { workspace = true } + +[dev-dependencies] +datafusion-expr = { workspace = true } +datafusion-functions-nested = { workspace = true } +insta = { workspace = true } +itertools = { workspace = true }
diff --git a/native/vendor/datafusion-pruning/LICENSE.txt b/native/vendor/datafusion-pruning/LICENSE.txt new file mode 100644 index 0000000..d74c6b5 --- /dev/null +++ b/native/vendor/datafusion-pruning/LICENSE.txt
@@ -0,0 +1,212 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + +This project includes code from Apache Aurora. + +* dev/release/{release,changelog,release-candidate} are based on the scripts from + Apache Aurora + +Copyright: 2016 The Apache Software Foundation. +Home page: https://aurora.apache.org/ +License: http://www.apache.org/licenses/LICENSE-2.0
diff --git a/native/vendor/datafusion-pruning/NOTICE.txt b/native/vendor/datafusion-pruning/NOTICE.txt new file mode 100644 index 0000000..0bd2d52 --- /dev/null +++ b/native/vendor/datafusion-pruning/NOTICE.txt
@@ -0,0 +1,5 @@ +Apache DataFusion +Copyright 2019-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/).
diff --git a/native/vendor/datafusion-pruning/README.md b/native/vendor/datafusion-pruning/README.md new file mode 100644 index 0000000..4db5091 --- /dev/null +++ b/native/vendor/datafusion-pruning/README.md
@@ -0,0 +1,34 @@ +<!--- + 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. +--> + +# Apache DataFusion Pruning Logic + +[Apache DataFusion] is an extensible query execution framework, written in Rust, that uses [Apache Arrow] as its in-memory format. + +This crate is a submodule of DataFusion that contains pruning logic, to analyze filter expressions with +statistics such as min/max values and null counts, proving files / large subsections of files can be skipped +without reading the actual data. + +Most projects should use the [`datafusion`] crate directly, which re-exports +this module. If you are already using the [`datafusion`] crate, there is no +reason to use this crate directly in your project as well. + +[apache arrow]: https://arrow.apache.org/ +[apache datafusion]: https://datafusion.apache.org/ +[`datafusion`]: https://crates.io/crates/datafusion
diff --git a/native/vendor/datafusion-pruning/src/file_pruner.rs b/native/vendor/datafusion-pruning/src/file_pruner.rs new file mode 100644 index 0000000..6618329 --- /dev/null +++ b/native/vendor/datafusion-pruning/src/file_pruner.rs
@@ -0,0 +1,171 @@ +// 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. + +//! File-level pruning based on partition values and file-level statistics + +use std::sync::Arc; + +use arrow::datatypes::{FieldRef, SchemaRef}; +use datafusion_common::{Result, internal_datafusion_err, pruning::PrunableStatistics}; +use datafusion_datasource::PartitionedFile; +use datafusion_physical_expr::DynamicFilterTracking; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_plan::metrics::Count; +use log::debug; + +use crate::build_pruning_predicate; + +/// Prune based on file-level statistics. +/// +/// Note: Partition column pruning is handled earlier via `replace_columns_with_literals` +/// which substitutes partition column references with their literal values before +/// the predicate reaches this pruner. +pub struct FilePruner { + predicate: Arc<dyn PhysicalExpr>, + /// Tracks the dynamic filters inside `predicate` so we only rebuild the + /// pruning predicate when one of them has actually moved. + tracking: DynamicFilterTracking, + /// Whether [`Self::should_prune`] has built+evaluated the pruning predicate + /// at least once. The first check always runs; subsequent checks only run + /// when a watched dynamic filter changed. + checked_once: bool, + /// Schema used for pruning (the logical file schema). + file_schema: SchemaRef, + file_stats_pruning: PrunableStatistics, + predicate_creation_errors: Count, +} + +impl FilePruner { + #[deprecated( + since = "52.0.0", + note = "Use `try_new` instead which returns None if no statistics are available" + )] + #[expect(clippy::needless_pass_by_value)] + pub fn new( + predicate: Arc<dyn PhysicalExpr>, + logical_file_schema: &SchemaRef, + _partition_fields: Vec<FieldRef>, + partitioned_file: PartitionedFile, + predicate_creation_errors: Count, + ) -> Result<Self> { + Self::try_new( + predicate, + logical_file_schema, + &partitioned_file, + predicate_creation_errors, + ) + .ok_or_else(|| { + internal_datafusion_err!( + "FilePruner::new called on a file without statistics: {:?}", + partitioned_file + ) + }) + } + + /// Create a file pruner for this file, or `None` when pruning it cannot + /// help. + /// + /// Returns `None` when the file has no statistics struct to evaluate a + /// pruning predicate against, or when the predicate is purely static and the + /// file has no usable column statistics — in that case planning already did + /// everything such a pruner could. A predicate carrying a dynamic filter is + /// always accepted (given a statistics struct), since it may prune via + /// partition-value folding even without column statistics. + pub fn try_new( + predicate: Arc<dyn PhysicalExpr>, + file_schema: &SchemaRef, + partitioned_file: &PartitionedFile, + predicate_creation_errors: Count, + ) -> Option<Self> { + // A pruning predicate is evaluated against a statistics struct, so one + // must exist (its columns may all be `Absent`). + let file_stats = partitioned_file.statistics.as_ref()?; + let tracking = DynamicFilterTracking::classify(&predicate); + // Only build a pruner when it could prune something planning didn't + // already: the file has real column statistics, or the predicate carries + // a dynamic filter (whose value, or folded partition columns, can prune + // even without column statistics). For a purely static predicate with no + // usable stats there is nothing to gain. + if !partitioned_file.has_statistics() && !tracking.contains_dynamic_filter() { + return None; + } + let file_stats_pruning = + PrunableStatistics::new(vec![file_stats.clone()], Arc::clone(file_schema)); + Some(Self { + predicate, + tracking, + checked_once: false, + file_schema: Arc::clone(file_schema), + file_stats_pruning, + predicate_creation_errors, + }) + } + + /// Returns `true` if this pruner watches a dynamic filter that can still + /// change, meaning [`Self::should_prune`] is worth re-checking as the scan + /// progresses. When `false`, the predicate is effectively static for the + /// remainder of the scan and the caller can avoid wrapping the stream in a + /// per-batch re-pruning adapter. + pub fn is_watching(&self) -> bool { + matches!(self.tracking, DynamicFilterTracking::Watching(_)) + } + + pub fn should_prune(&mut self) -> Result<bool> { + // Building the pruning predicate is expensive (it involves expression + // analysis), so we only do it on the first check and whenever a dynamic + // filter inside the predicate has actually moved. + // + // Dynamic filter expressions can change their values during query + // execution; `DynamicFilterTracking` watches the still-incomplete + // filters and reports a change at most once per update. A purely static + // predicate (or one whose dynamic filters have all completed) is checked + // exactly once. + let should_build = if self.checked_once { + self.tracking.watcher().is_some_and(|w| w.changed()) + } else { + self.checked_once = true; + true + }; + if !should_build { + return Ok(false); + } + let pruning_predicate = build_pruning_predicate( + Arc::clone(&self.predicate), + &self.file_schema, + &self.predicate_creation_errors, + ); + let Some(pruning_predicate) = pruning_predicate else { + return Ok(false); + }; + match pruning_predicate.prune(&self.file_stats_pruning) { + Ok(values) => { + assert!(values.len() == 1); + // We expect a single container -> if all containers are false skip this file + if values.into_iter().all(|v| !v) { + return Ok(true); + } + } + // Stats filter array could not be built, so we can't prune + Err(e) => { + debug!("Ignoring error building pruning predicate for file: {e}"); + self.predicate_creation_errors.add(1); + } + } + + Ok(false) + } +}
diff --git a/native/vendor/datafusion-pruning/src/key_range_bitmap_expr.rs b/native/vendor/datafusion-pruning/src/key_range_bitmap_expr.rs new file mode 100644 index 0000000..ba30a05 --- /dev/null +++ b/native/vendor/datafusion-pruning/src/key_range_bitmap_expr.rs
@@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt::{self, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef}; +use datafusion_physical_plan::ColumnarValue; +use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; + +/// Tests containers' `[min, max]` statistics against a build side's +/// [`KeyRangeBitmap`]: one nullable Boolean per container, where `false` proves +/// no build key lies in it and `NULL` means the bounds could not decide. +#[derive(Debug)] +pub(crate) struct KeyRangeBitmapPruningExpr { + pub(crate) min: PhysicalExprRef, + pub(crate) max: PhysicalExprRef, + pub(crate) bitmap: Arc<KeyRangeBitmap>, +} + +impl PartialEq for KeyRangeBitmapPruningExpr { + fn eq(&self, other: &Self) -> bool { + self.min.eq(&other.min) + && self.max.eq(&other.max) + && Arc::ptr_eq(&self.bitmap, &other.bitmap) + } +} +impl Eq for KeyRangeBitmapPruningExpr {} + +impl Hash for KeyRangeBitmapPruningExpr { + fn hash<H: Hasher>(&self, state: &mut H) { + self.min.hash(state); + self.max.hash(state); + Arc::as_ptr(&self.bitmap).hash(state); + } +} + +impl Display for KeyRangeBitmapPruningExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let (set, total) = self.bitmap.fill_stats(); + write!( + f, + "KEY_RANGE_BITMAP({}, {}, {set}/{total} buckets)", + self.min, self.max + ) + } +} + +impl PhysicalExpr for KeyRangeBitmapPruningExpr { + fn data_type(&self, _input_schema: &Schema) -> Result<DataType> { + Ok(DataType::Boolean) + } + + fn nullable(&self, _input_schema: &Schema) -> Result<bool> { + Ok(true) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + let rows = batch.num_rows(); + let min = self.min.evaluate(batch)?.into_array(rows)?; + let max = self.max.evaluate(batch)?.into_array(rows)?; + let matches = self.bitmap.may_contain_ranges(&min, &max); + Ok(ColumnarValue::Array(Arc::new(matches))) + } + + fn children(&self) -> Vec<&PhysicalExprRef> { + vec![&self.min, &self.max] + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<PhysicalExprRef>, + ) -> Result<PhysicalExprRef> { + assert_eq_or_internal_err!(children.len(), 2); + Ok(Arc::new(Self { + min: Arc::clone(&children[0]), + max: Arc::clone(&children[1]), + bitmap: Arc::clone(&self.bitmap), + })) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{self}") + } +}
diff --git a/native/vendor/datafusion-pruning/src/lib.rs b/native/vendor/datafusion-pruning/src/lib.rs new file mode 100644 index 0000000..7bde7e2 --- /dev/null +++ b/native/vendor/datafusion-pruning/src/lib.rs
@@ -0,0 +1,28 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![cfg_attr(test, allow(clippy::needless_pass_by_value))] + +mod file_pruner; +mod key_range_bitmap_expr; +mod pruning_predicate; + +pub use file_pruner::FilePruner; +pub use pruning_predicate::{ + MAX_IN_LIST_SIZE, PredicateRewriter, PruningPredicate, PruningPredicateBuilder, + PruningStatistics, RequiredColumns, UnhandledPredicateHook, build_pruning_predicate, +};
diff --git a/native/vendor/datafusion-pruning/src/pruning_predicate.rs b/native/vendor/datafusion-pruning/src/pruning_predicate.rs new file mode 100644 index 0000000..812cc78 --- /dev/null +++ b/native/vendor/datafusion-pruning/src/pruning_predicate.rs
@@ -0,0 +1,6424 @@ +// 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. + +//! [`PruningPredicate`] to apply filter [`Expr`] to prune "containers" +//! based on statistics (e.g. Parquet Row Groups) +//! +//! [`Expr`]: https://docs.rs/datafusion/latest/datafusion/logical_expr/enum.Expr.html +use std::collections::HashSet; +use std::sync::Arc; + +use crate::key_range_bitmap_expr::KeyRangeBitmapPruningExpr; + +use arrow::array::AsArray; +use arrow::{ + array::{ArrayRef, BooleanArray, new_null_array}, + datatypes::{DataType, Field, Schema, SchemaRef}, + record_batch::{RecordBatch, RecordBatchOptions}, +}; +// pub use for backwards compatibility +pub use datafusion_common::pruning::PruningStatistics; +use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; +use datafusion_physical_plan::metrics::Count; +use log::{debug, trace}; + +use datafusion_common::error::Result; +use datafusion_common::tree_node::{TransformedResult, TreeNodeRecursion}; +use datafusion_common::{ + _internal_datafusion_err, Column, DFSchema, assert_eq_or_internal_err, +}; +use datafusion_common::{ + ScalarValue, internal_datafusion_err, plan_datafusion_err, plan_err, + tree_node::{Transformed, TreeNode}, +}; +use datafusion_expr_common::casts::try_cast_literal_to_type; +use datafusion_expr_common::operator::Operator; +use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee}; +use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr}; +use datafusion_physical_expr_common::physical_expr::snapshot_physical_expr_opt; +use datafusion_physical_plan::joins::HashTableLookupExpr; +use datafusion_physical_plan::{ColumnarValue, PhysicalExpr}; + +/// Used to prove that arbitrary predicates (boolean expression) can not +/// possibly evaluate to `true` given information about a column provided by +/// [`PruningStatistics`]. +/// +/// # Introduction +/// +/// `PruningPredicate` analyzes filter expressions using statistics such as +/// min/max values and null counts, attempting to prove a "container" (e.g. +/// Parquet Row Group) can be skipped without reading the actual data, +/// potentially leading to significant performance improvements. +/// +/// For example, `PruningPredicate`s are used to prune Parquet Row Groups based +/// on the min/max values found in the Parquet metadata. If the +/// `PruningPredicate` can prove that the filter can never evaluate to `true` +/// for any row in the Row Group, the entire Row Group is skipped during query +/// execution. +/// +/// The `PruningPredicate` API is general, and can be used for pruning other +/// types of containers (e.g. files) based on statistics that may be known from +/// external catalogs (e.g. Delta Lake) or other sources. How this works is a +/// subtle topic. See the Background and Implementation section for details. +/// +/// `PruningPredicate` supports: +/// +/// 1. Arbitrary expressions (including user defined functions) +/// +/// 2. Vectorized evaluation (provide more than one set of statistics at a time) +/// so it is suitable for pruning 1000s of containers. +/// +/// 3. Any source of information that implements the [`PruningStatistics`] trait +/// (not just Parquet metadata). +/// +/// # Example +/// +/// See the [`pruning.rs` example in the `datafusion-examples`] for a complete +/// example of how to use `PruningPredicate` to prune files based on min/max +/// values. +/// +/// [`pruning.rs` example in the `datafusion-examples`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/query_planning/pruning.rs +/// +/// Given an expression like `x = 5` and statistics for 3 containers (Row +/// Groups, files, etc) `A`, `B`, and `C`: +/// +/// ```text +/// A: {x_min = 0, x_max = 4} +/// B: {x_min = 2, x_max = 10} +/// C: {x_min = 5, x_max = 8} +/// ``` +/// +/// `PruningPredicate` will conclude that the rows in container `A` can never +/// be true (as the maximum value is only `4`), so it can be pruned: +/// +/// ```text +/// A: false (no rows could possibly match x = 5) +/// B: true (rows might match x = 5) +/// C: true (rows might match x = 5) +/// ``` +/// +/// See [`PruningPredicateBuilder`] and [`PruningPredicate::prune`] for more information. +/// +/// # Background +/// +/// ## Boolean Tri-state logic +/// +/// To understand the details of the rest of this documentation, it is important +/// to understand how the tri-state boolean logic in SQL works. As this is +/// somewhat esoteric, we review it here. +/// +/// SQL has a notion of `NULL` that represents the value is `“unknown”` and this +/// uncertainty propagates through expressions. SQL `NULL` behaves very +/// differently than the `NULL` in most other languages where it is a special, +/// sentinel value (e.g. `0` in `C/C++`). While representing uncertainty with +/// `NULL` is powerful and elegant, SQL `NULL`s are often deeply confusing when +/// first encountered as they behave differently than most programmers may +/// expect. +/// +/// In most other programming languages, +/// * `a == NULL` evaluates to `true` if `a` also had the value `NULL` +/// * `a == NULL` evaluates to `false` if `a` has any other value +/// +/// However, in SQL `a = NULL` **always** evaluates to `NULL` (never `true` or +/// `false`): +/// +/// Expression | Result +/// ------------- | --------- +/// `1 = NULL` | `NULL` +/// `NULL = NULL` | `NULL` +/// +/// Also important is how `AND` and `OR` works with tri-state boolean logic as +/// (perhaps counterintuitively) the result is **not** always NULL. While +/// consistent with the notion of `NULL` representing “unknown”, this is again, +/// often deeply confusing 🤯 when first encountered. +/// +/// Expression | Result | Intuition +/// --------------- | --------- | ----------- +/// `NULL AND true` | `NULL` | The `NULL` stands for “unknown” and if it were `true` or `false` the overall expression value could change +/// `NULL AND false` | `false` | If the `NULL` was either `true` or `false` the overall expression is still `false` +/// `NULL AND NULL` | `NULL` | +/// +/// Expression | Result | Intuition +/// --------------- | --------- | ---------- +/// `NULL OR true` | `true` | If the `NULL` was either `true` or `false` the overall expression is still `true` +/// `NULL OR false` | `NULL` | The `NULL` stands for “unknown” and if it were `true` or `false` the overall expression value could change +/// `NULL OR NULL` | `NULL` | +/// +/// ## SQL Filter Semantics +/// +/// The SQL `WHERE` clause has a boolean expression, often called a filter or +/// predicate. The semantics of this predicate are that the query evaluates the +/// predicate for each row in the input tables and: +/// +/// * Rows that evaluate to `true` are returned in the query results +/// +/// * Rows that evaluate to `false` are not returned (“filtered out” or “pruned” or “skipped”). +/// +/// * Rows that evaluate to `NULL` are **NOT** returned (also “filtered out”). +/// Note: *this treatment of `NULL` is **DIFFERENT** than how `NULL` is treated +/// in the rewritten predicate described below.* +/// +/// # `PruningPredicate` Implementation +/// +/// Armed with the information in the Background section, we can now understand +/// how the `PruningPredicate` logic works. +/// +/// ## Interface +/// +/// **Inputs** +/// 1. An input schema describing what columns exist +/// +/// 2. A predicate (expression that evaluates to a boolean) +/// +/// 3. [`PruningStatistics`] that provides information about columns in that +/// schema, for multiple “containers”. For each column in each container, it +/// provides optional information on contained values, min_values, max_values, +/// null_counts counts, and row_counts counts. +/// +/// **Outputs**: +/// A (non null) boolean value for each container: +/// * `true`: There MAY be rows that match the predicate +/// +/// * `false`: There are no rows that could possibly match the predicate (the +/// predicate can never possibly be true). The container can be pruned (skipped) +/// entirely. +/// +/// While `PruningPredicate` will never return a `NULL` value, the +/// rewritten predicate (as returned by `build_predicate_expression` and used internally +/// by `PruningPredicate`) may evaluate to `NULL` when some of the min/max values +/// or null / row counts are not known. +/// +/// In order to be correct, `PruningPredicate` must return false +/// **only** if it can determine that for all rows in the container, the +/// predicate could never evaluate to `true` (always evaluates to either `NULL` +/// or `false`). +/// +/// ## Contains Analysis and Min/Max Rewrite +/// +/// `PruningPredicate` works by first analyzing the predicate to see what +/// [`LiteralGuarantee`] must hold for the predicate to be true. +/// +/// Then, the `PruningPredicate` rewrites the original predicate into an +/// expression that references the min/max values of each column in the original +/// predicate. +/// +/// When the min/max values are actually substituted in to this expression and +/// evaluated, the result means +/// +/// * `true`: there MAY be rows that pass the predicate, **KEEPS** the container +/// +/// * `NULL`: there MAY be rows that pass the predicate, **KEEPS** the container +/// Note that rewritten predicate can evaluate to NULL when some of +/// the min/max values are not known. *Note that this is different than +/// the SQL filter semantics where `NULL` means the row is filtered +/// out.* +/// +/// * `false`: there are no rows that could possibly match the predicate, +/// **PRUNES** the container +/// +/// For example, given a column `x`, the `x_min`, `x_max`, `x_null_count`, and +/// `x_row_count` represent the minimum and maximum values, the null count of +/// column `x`, and the row count of column `x`, provided by the `PruningStatistics`. +/// `x_null_count` and `x_row_count` are used to handle the case where the column `x` +/// is known to be all `NULL`s. Note this is different from knowing nothing about +/// the column `x`, which confusingly is encoded by returning `NULL` for the min/max +/// values from [`PruningStatistics::max_values`] and [`PruningStatistics::min_values`]. +/// +/// Here are some examples of the rewritten predicates: +/// +/// Original Predicate | Rewritten Predicate +/// ------------------ | -------------------- +/// `x = 5` | `x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max)` +/// `x < 5` | `x_null_count != x_row_count AND (x_min < 5)` +/// `x = 5 AND y = 10` | `x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max) AND y_null_count != y_row_count (y_min <= 10 AND 10 <= y_max)` +/// `x IS NULL` | `x_null_count > 0` +/// `x IS NOT NULL` | `x_null_count != row_count` +/// `CAST(x as int) = 5` | `x_null_count != x_row_count (CAST(x_min as int) <= 5 AND 5 <= CAST(x_max as int))` +/// +/// ## Predicate Evaluation +/// The PruningPredicate works in two passes +/// +/// **First pass**: For each `LiteralGuarantee` calls +/// [`PruningStatistics::contained`] and rules out containers where the +/// LiteralGuarantees are not satisfied +/// +/// **Second Pass**: Evaluates the rewritten expression using the +/// min/max/null_counts/row_counts values for each column for each container. For any +/// container that this expression evaluates to `false`, it rules out those +/// containers. +/// +/// +/// ### Example 1 +/// +/// Given the predicate, `x = 5 AND y = 10`, the rewritten predicate would look like: +/// +/// ```sql +/// x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max) +/// AND +/// y_null_count != y_row_count AND (y_min <= 10 AND 10 <= y_max) +/// ``` +/// +/// If we know that for a given container, `x` is between `1 and 100` and we know that +/// `y` is between `4` and `7`, we know nothing about the null count and row count of +/// `x` and `y`, the input statistics might look like: +/// +/// Column | Value +/// -------- | ----- +/// `x_min` | `1` +/// `x_max` | `100` +/// `x_null_count` | `null` +/// `x_row_count` | `null` +/// `y_min` | `4` +/// `y_max` | `7` +/// `y_null_count` | `null` +/// `y_row_count` | `null` +/// +/// When these statistics values are substituted in to the rewritten predicate and +/// simplified, the result is `false`: +/// +/// * `null != null AND (1 <= 5 AND 5 <= 100) AND null != null AND (4 <= 10 AND 10 <= 7)` +/// * `null = null` is `null` which is not true, so the AND moves on to the next clause +/// * `null and (1 <= 5 AND 5 <= 100) AND null AND (4 <= 10 AND 10 <= 7)` +/// * evaluating the clauses further we get: +/// * `null and true and null and false` +/// * `null and false` +/// * `false` +/// +/// Returning `false` means the container can be pruned, which matches the +/// intuition that `x = 5 AND y = 10` can’t be true for any row if all values of `y` +/// are `7` or less. +/// +/// Note that if we had ended up with `null AND true AND null AND true` the result +/// would have been `null`. +/// `null` is treated the same as`true`, because we can't prove that the predicate is `false.` +/// +/// If, for some other container, we knew `y` was between the values `4` and +/// `15`, then the rewritten predicate evaluates to `true` (verifying this is +/// left as an exercise to the reader -- are you still here?), and the container +/// **could not** be pruned. The intuition is that there may be rows where the +/// predicate *might* evaluate to `true`, and the only way to find out is to do +/// more analysis, for example by actually reading the data and evaluating the +/// predicate row by row. +/// +/// ### Example 2 +/// +/// Given the same predicate, `x = 5 AND y = 10`, the rewritten predicate would +/// look like the same as example 1: +/// +/// ```sql +/// x_null_count != x_row_count AND (x_min <= 5 AND 5 <= x_max) +/// AND +/// y_null_count != y_row_count AND (y_min <= 10 AND 10 <= y_max) +/// ``` +/// +/// If we know that for another given container, `x_min` is NULL and `x_max` is +/// NULL (the min/max values are unknown), `x_null_count` is `100` and `x_row_count` +/// is `100`; we know that `y` is between `4` and `7`, but we know nothing about +/// the null count and row count of `y`. The input statistics might look like: +/// +/// Column | Value +/// -------- | ----- +/// `x_min` | `null` +/// `x_max` | `null` +/// `x_null_count` | `100` +/// `x_row_count` | `100` +/// `y_min` | `4` +/// `y_max` | `7` +/// `y_null_count` | `null` +/// `y_row_count` | `null` +/// +/// When these statistics values are substituted in to the rewritten predicate and +/// simplified, the result is `false`: +/// +/// * `100 != 100 AND (null <= 5 AND 5 <= null) AND null = null AND (4 <= 10 AND 10 <= 7)` +/// * `false AND null AND null AND false` +/// * `false AND false` +/// * `false` +/// +/// Returning `false` means the container can be pruned, which matches the +/// intuition that `x = 5 AND y = 10` can’t be true because all values in `x` +/// are known to be NULL. +/// +/// # Related Work +/// +/// [`PruningPredicate`] implements the type of min/max pruning described in +/// Section `3.3.3` of the [`Snowflake SIGMOD Paper`]. The technique is +/// described by various research such as [small materialized aggregates], [zone +/// maps], and [data skipping]. +/// +/// [`Snowflake SIGMOD Paper`]: https://dl.acm.org/doi/10.1145/2882903.2903741 +/// [small materialized aggregates]: https://www.vldb.org/conf/1998/p476.pdf +/// [zone maps]: https://dl.acm.org/doi/10.1007/978-3-642-03730-6_10 +/// [data skipping]: https://dl.acm.org/doi/10.1145/2588555.2610515 +#[derive(Debug, Clone)] +pub struct PruningPredicate { + /// The input schema against which the predicate will be evaluated + schema: SchemaRef, + /// A min/max pruning predicate (rewritten in terms of column min/max + /// values, which are supplied by statistics) + predicate_expr: Arc<dyn PhysicalExpr>, + /// Description of which statistics are required to evaluate `predicate_expr` + required_columns: RequiredColumns, + /// Original physical predicate from which this predicate expr is derived + /// (required for serialization) + orig_expr: Arc<dyn PhysicalExpr>, + /// [`LiteralGuarantee`]s used to try and prove a predicate can not possibly + /// evaluate to `true`. + /// + /// See [`PruningPredicate::literal_guarantees`] for more details. + literal_guarantees: Vec<LiteralGuarantee>, +} + +/// Build a pruning predicate from an optional predicate expression. +/// If the predicate is None or the predicate cannot be converted to a pruning +/// predicate, return None. +/// If there is an error creating the pruning predicate it is recorded by incrementing +/// the `predicate_creation_errors` counter. +pub fn build_pruning_predicate( + predicate: Arc<dyn PhysicalExpr>, + file_schema: &SchemaRef, + predicate_creation_errors: &Count, +) -> Option<Arc<PruningPredicate>> { + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(file_schema)) + .with_error_counter(predicate_creation_errors) + .build(predicate) +} + +/// Builder for a [`PruningPredicate`]. Groups optional configuration — +/// `IN (...)` rewrite cap, error counter — so future additions do not +/// churn the top-level API. +/// +/// The two entry points are: +/// - [`Self::build`]: convenience for scan sites that already track a +/// `predicate_creation_errors` counter. Returns `Some(Arc<..>)` when the +/// resulting predicate can actually prune, `None` when it is trivially +/// true or when construction failed (in which case the error counter is +/// incremented if one was supplied). +/// - [`Self::try_build`]: returns a raw `Result<PruningPredicate>` for +/// callers that want to surface errors themselves. +/// +#[derive(Default)] +pub struct PruningPredicateBuilder<'a> { + file_schema: Option<SchemaRef>, + error_counter: Option<&'a Count>, + max_in_list_size: usize, +} + +impl<'a> PruningPredicateBuilder<'a> { + /// Create a new builder with the default pruning predicate configuration. + pub fn new() -> Self { + Self { + file_schema: None, + error_counter: None, + max_in_list_size: MAX_IN_LIST_SIZE, + } + } + + /// Set the schema of the container that will be pruned (typically the + /// parquet file schema). + pub fn with_file_schema(mut self, file_schema: SchemaRef) -> Self { + self.file_schema = Some(file_schema); + self + } + + /// Metric counter incremented once per predicate that fails to build. + /// Only consulted by [`Self::build`]; [`Self::try_build`] surfaces the + /// error directly. + pub fn with_error_counter(mut self, error_counter: &'a Count) -> Self { + self.error_counter = Some(error_counter); + self + } + + /// Cap on the size of `IN (...)` lists that will be rewritten into per- + /// value min/max statistics checks. Lists longer than this fall back to + /// the unhandled-predicate hook (typically "keep the container"). + /// + /// Query engines typically pass + /// `datafusion.execution.parquet.max_in_list_size` here. + pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self { + self.max_in_list_size = max_in_list_size; + self + } + + /// Build a [`PruningPredicate`] wrapped in `Some(Arc<..>)` when it can + /// prune, `None` when it is trivially true or when construction fails. + /// If [`Self::with_error_counter`] was set, construction failures are + /// recorded there. + pub fn build( + self, + predicate: Arc<dyn PhysicalExpr>, + ) -> Option<Arc<PruningPredicate>> { + let error_counter = self.error_counter; + match self.try_build(predicate) { + Ok(pruning_predicate) => { + if !pruning_predicate.always_true() { + return Some(Arc::new(pruning_predicate)); + } + } + Err(e) => { + debug!("Could not create pruning predicate for: {e}"); + if let Some(counter) = error_counter { + counter.add(1); + } + } + } + None + } + + /// Build a [`PruningPredicate`], returning the construction error + /// directly. Callers that want the always-true predicate elided or + /// errors folded into a counter should use [`Self::build`] instead. + pub fn try_build( + self, + mut predicate: Arc<dyn PhysicalExpr>, + ) -> Result<PruningPredicate> { + let file_schema = self.file_schema.ok_or_else(|| { + _internal_datafusion_err!( + "PruningPredicateBuilder requires a file schema (call `with_file_schema`)" + ) + })?; + + // Get a (simpler) snapshot of the physical expr here to use with `PruningPredicate`. + // In particular this unravels any `DynamicFilterPhysicalExpr`s by snapshotting them + // so that PruningPredicate can work with a static expression. + let tf = snapshot_physical_expr_opt(predicate)?; + if tf.transformed { + // If we had an expression such as Dynamic(part_col < 5 and col < 10) + // (this could come from something like `select * from t order by part_col, col, limit 10`) + // after snapshotting and because `DynamicFilterPhysicalExpr` applies child replacements to its + // children after snapshotting and previously `replace_columns_with_literals` may have been called with partition values + // the expression we have now is `8 < 5 and col < 10`. + // Thus we need as simplifier pass to get `false and col < 10` => `false` here. + let simplifier = PhysicalExprSimplifier::new(&file_schema); + predicate = simplifier.simplify(tf.data)?; + } else { + predicate = tf.data; + } + let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _; + + // build predicate expression once + let mut required_columns = RequiredColumns::new(); + let predicate_expr = build_predicate_expression( + &predicate, + &file_schema, + &mut required_columns, + &unhandled_hook, + self.max_in_list_size, + ); + let predicate_schema = required_columns.schema(); + // Simplify the newly created predicate to get rid of redundant casts, comparisons, etc. + let predicate_expr = + PhysicalExprSimplifier::new(&predicate_schema).simplify(predicate_expr)?; + let literal_guarantees = LiteralGuarantee::analyze(&predicate); + + Ok(PruningPredicate { + schema: file_schema, + predicate_expr, + required_columns, + orig_expr: predicate, + literal_guarantees, + }) + } +} + +/// Rewrites predicates that [`PredicateRewriter`] can not handle, e.g. certain +/// complex expressions or predicates that reference columns that are not in the +/// schema. +pub trait UnhandledPredicateHook { + /// Called when a predicate can not be rewritten in terms of statistics or + /// references a column that is not in the schema. + fn handle(&self, expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr>; +} + +/// The default handling for unhandled predicates is to return a constant `true` +/// (meaning don't prune the container) +#[derive(Debug, Clone)] +struct ConstantUnhandledPredicateHook { + default: Arc<dyn PhysicalExpr>, +} + +impl Default for ConstantUnhandledPredicateHook { + fn default() -> Self { + Self { + default: Arc::new(phys_expr::Literal::new(ScalarValue::from(true))), + } + } +} + +impl UnhandledPredicateHook for ConstantUnhandledPredicateHook { + fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> { + Arc::clone(&self.default) + } +} + +impl PruningPredicate { + /// Try to create a new instance of [`PruningPredicate`] + /// + /// This will translate the provided `expr` filter expression into + /// a *pruning predicate*. + /// + /// A pruning predicate is one that has been rewritten in terms of + /// the min and max values of column references and that evaluates + /// to FALSE if the filter predicate would evaluate FALSE *for + /// every row* whose values fell within the min / max ranges (aka + /// could be pruned). + /// + /// The pruning predicate evaluates to TRUE or NULL + /// if the filter predicate *might* evaluate to TRUE for at least + /// one row whose values fell within the min/max ranges (in other + /// words they might pass the predicate) + /// + /// For example, the filter expression `(column / 2) = 4` becomes + /// the pruning predicate + /// `(column_min / 2) <= 4 && 4 <= (column_max / 2))` + /// + /// See the struct level documentation on [`PruningPredicate`] for more + /// details. + /// + /// Note that `PruningPredicate` does not attempt to normalize or simplify + /// the input expression unless calling [`snapshot_physical_expr_opt`] + /// returns a new expression. + /// It is recommended that you pass the expressions through [`PhysicalExprSimplifier`] + /// before calling this method to make sure the expressions can be used for pruning. + /// + /// Use [`PruningPredicateBuilder`] to construct new pruning predicates. + #[deprecated(since = "55.0.0", note = "Use PruningPredicateBuilder instead")] + pub fn try_new(expr: Arc<dyn PhysicalExpr>, schema: SchemaRef) -> Result<Self> { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr) + } + + /// For each set of statistics, evaluates the pruning predicate + /// and returns a `bool` with the following meaning for a + /// all rows whose values match the statistics: + /// + /// `true`: There MAY be rows that match the predicate + /// + /// `false`: There are no rows that could possibly match the predicate + /// + /// Note: the predicate passed to `prune` should already be simplified as + /// much as possible (e.g. this pass doesn't handle some + /// expressions like `b = false`, but it does handle the + /// simplified version `b`. See [`ExprSimplifier`] to simplify expressions. + /// + /// [`ExprSimplifier`]: https://docs.rs/datafusion/latest/datafusion/optimizer/simplify_expressions/struct.ExprSimplifier.html + pub fn prune<S: PruningStatistics + ?Sized>( + &self, + statistics: &S, + ) -> Result<Vec<bool>> { + let mut builder = BoolVecBuilder::new(statistics.num_containers()); + + // Try to prove the predicate can't be true for the containers based on + // literal guarantees + for literal_guarantee in &self.literal_guarantees { + let LiteralGuarantee { + column, + guarantee, + literals, + } = literal_guarantee; + if let Some(results) = statistics.contained(column, literals) { + match guarantee { + // `In` means the values in the column must be one of the + // values in the set for the predicate to evaluate to true. + // If `contained` returns false, that means the column is + // not any of the values so we can prune the container + Guarantee::In => builder.combine_array(&results), + // `NotIn` means the values in the column must not be + // any of the values in the set for the predicate to + // evaluate to true. If `contained` returns true, it means the + // column is only in the set of values so we can prune the + // container + Guarantee::NotIn => { + builder.combine_array(&arrow::compute::not(&results)?) + } + } + // if all containers are pruned (has rows that DEFINITELY DO NOT pass the predicate) + // can return early without evaluating the rest of predicates. + if builder.check_all_pruned() { + return Ok(builder.build()); + } + } + } + + // Next, try to prove the predicate can't be true for the containers based + // on min/max values + + // build a RecordBatch that contains the min/max values in the + // appropriate statistics columns for the min/max predicate + let statistics_batch = + build_statistics_record_batch(statistics, &self.required_columns)?; + + // Evaluate the pruning predicate on that record batch and append any results to the builder + builder.combine_value(self.predicate_expr.evaluate(&statistics_batch)?); + + Ok(builder.build()) + } + + /// Return a reference to the input schema + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Returns a reference to the physical expr used to construct this pruning predicate + pub fn orig_expr(&self) -> &Arc<dyn PhysicalExpr> { + &self.orig_expr + } + + /// Returns a reference to the predicate expr + pub fn predicate_expr(&self) -> &Arc<dyn PhysicalExpr> { + &self.predicate_expr + } + + /// Returns a reference to the literal guarantees + /// + /// Note that **All** `LiteralGuarantee`s must be satisfied for the + /// expression to possibly be `true`. If any is not satisfied, the + /// expression is guaranteed to be `null` or `false`. + pub fn literal_guarantees(&self) -> &[LiteralGuarantee] { + &self.literal_guarantees + } + + /// Returns true if this pruning predicate can not prune anything. + /// + /// This happens if the predicate is a literal `true` and + /// literal_guarantees is empty. + /// + /// This can happen when a predicate is simplified to a constant `true` + pub fn always_true(&self) -> bool { + is_always_true(&self.predicate_expr) && self.literal_guarantees.is_empty() + } + + pub fn required_columns(&self) -> &RequiredColumns { + &self.required_columns + } + + /// Names of the columns that are known to be / not be in a set + /// of literals (constants). These are the columns the that may be passed to + /// [`PruningStatistics::contained`] during pruning. + /// + /// This is useful to avoid fetching statistics for columns that will not be + /// used in the predicate. For example, it can be used to avoid reading + /// unneeded bloom filters (a non trivial operation). + pub fn literal_columns(&self) -> Vec<String> { + let mut seen = HashSet::new(); + self.literal_guarantees + .iter() + .map(|e| &e.column.name) + // avoid duplicates + .filter(|name| seen.insert(*name)) + .map(|s| s.to_string()) + .collect() + } +} + +/// Builds the return `Vec` for [`PruningPredicate::prune`]. +#[derive(Debug)] +struct BoolVecBuilder { + /// One element per container. Each element is + /// * `true`: if the container has row that may pass the predicate + /// * `false`: if the container has rows that DEFINITELY DO NOT pass the predicate + inner: Vec<bool>, +} + +impl BoolVecBuilder { + /// Create a new `BoolVecBuilder` with `num_containers` elements + fn new(num_containers: usize) -> Self { + Self { + // assume by default all containers may pass the predicate + inner: vec![true; num_containers], + } + } + + /// Combines result `array` for a conjunct (e.g. `AND` clause) of a + /// predicate into the currently in progress array. + /// + /// Each `array` element is: + /// * `true`: container has row that may pass the predicate + /// * `false`: all container rows DEFINITELY DO NOT pass the predicate + /// * `null`: container may or may not have rows that pass the predicate + fn combine_array(&mut self, array: &BooleanArray) { + assert_eq!(array.len(), self.inner.len()); + for (cur, new) in self.inner.iter_mut().zip(array.iter()) { + // `false` for this conjunct means we know for sure no rows could + // pass the predicate and thus we set the corresponding container + // location to false. + if let Some(false) = new { + *cur = false; + } + } + } + + /// Combines the results in the [`ColumnarValue`] to the currently in + /// progress array, following the same rules as [`Self::combine_array`]. + /// + /// # Panics + /// If `value` is not boolean + fn combine_value(&mut self, value: ColumnarValue) { + match value { + ColumnarValue::Array(array) => { + self.combine_array(array.as_boolean()); + } + ColumnarValue::Scalar(ScalarValue::Boolean(Some(false))) => { + // False means all containers can not pass the predicate + self.inner = vec![false; self.inner.len()]; + } + _ => { + // Null or true means the rows in container may pass this + // conjunct so we can't prune any containers based on that + } + } + } + + /// Convert this builder into a Vec of bools + fn build(self) -> Vec<bool> { + self.inner + } + + /// Check all containers has rows that DEFINITELY DO NOT pass the predicate + fn check_all_pruned(&self) -> bool { + self.inner.iter().all(|&x| !x) + } +} + +fn is_always_true(expr: &Arc<dyn PhysicalExpr>) -> bool { + expr.downcast_ref::<phys_expr::Literal>() + .map(|l| matches!(l.value(), ScalarValue::Boolean(Some(true)))) + .unwrap_or_default() +} + +fn is_always_false(expr: &Arc<dyn PhysicalExpr>) -> bool { + expr.downcast_ref::<phys_expr::Literal>() + .map(|l| matches!(l.value(), ScalarValue::Boolean(Some(false)))) + .unwrap_or_default() +} + +/// Describes which columns statistics are necessary to evaluate a +/// [`PruningPredicate`]. +/// +/// This structure permits reading and creating the minimum number statistics, +/// which is important since statistics may be non trivial to read (e.g. large +/// strings or when there are 1000s of columns). +/// +/// Handles creating references to the min/max statistics +/// for columns as well as recording which statistics are needed +#[derive(Debug, Default, Clone)] +pub struct RequiredColumns { + /// The statistics required to evaluate this predicate: + /// * The unqualified column in the input schema + /// * Statistics type (e.g. Min or Max or Null_Count) + /// * The field the statistics value should be placed in for + /// pruning predicate evaluation (e.g. `min_value` or `max_value`) + columns: Vec<(phys_expr::Column, StatisticsType, Field)>, +} + +impl RequiredColumns { + fn new() -> Self { + Self::default() + } + + /// Returns Some(column) if this is a single column predicate. + /// + /// Returns None if this is a multi-column predicate. + /// + /// Examples: + /// * `a > 5 OR a < 10` returns `Some(a)` + /// * `a > 5 OR b < 10` returns `None` + /// * `true` returns None + pub fn single_column(&self) -> Option<&phys_expr::Column> { + if self.columns.windows(2).all(|w| { + // check if all columns are the same (ignoring statistics and field) + let c1 = &w[0].0; + let c2 = &w[1].0; + c1 == c2 + }) { + self.columns.first().map(|r| &r.0) + } else { + None + } + } + + /// Returns a schema that describes the columns required to evaluate this + /// pruning predicate. + /// The schema contains the fields for each column in `self.columns` with + /// the appropriate data type for the statistics. + /// Order matters, this same order is used to evaluate the + /// pruning predicate. + fn schema(&self) -> Schema { + let fields = self + .columns + .iter() + .map(|(_c, _t, f)| f.clone()) + .collect::<Vec<_>>(); + Schema::new(fields) + } + + /// Returns an iterator over items in columns (see doc on + /// `self.columns` for details) + pub(crate) fn iter( + &self, + ) -> impl Iterator<Item = &(phys_expr::Column, StatisticsType, Field)> { + self.columns.iter() + } + + fn find_stat_column( + &self, + column: &phys_expr::Column, + statistics_type: StatisticsType, + ) -> Option<usize> { + match statistics_type { + StatisticsType::RowCount => { + // Use the first row count we find, if any + self.columns + .iter() + .enumerate() + .find(|(_i, (_c, t, _f))| t == &statistics_type) + .map(|(i, (_c, _t, _f))| i) + } + _ => self + .columns + .iter() + .enumerate() + .find(|(_i, (c, t, _f))| c == column && t == &statistics_type) + .map(|(i, (_c, _t, _f))| i), + } + } + + /// Rewrites column_expr so that all appearances of column + /// are replaced with a reference to either the min or max + /// statistics column, while keeping track that a reference to the statistics + /// column is required + /// + /// for example, an expression like `col("foo") > 5`, when called + /// with Max would result in an expression like `col("foo_max") > + /// 5` with the appropriate entry noted in self.columns + fn stat_column_expr( + &mut self, + column: &phys_expr::Column, + column_expr: &Arc<dyn PhysicalExpr>, + field: &Field, + stat_type: StatisticsType, + ) -> Result<Arc<dyn PhysicalExpr>> { + let (idx, need_to_insert) = match self.find_stat_column(column, stat_type) { + Some(idx) => (idx, false), + None => (self.columns.len(), true), + }; + + let column_name = column.name(); + let stat_column_name = match stat_type { + StatisticsType::Min => format!("{column_name}_min"), + StatisticsType::Max => format!("{column_name}_max"), + StatisticsType::NullCount => format!("{column_name}_null_count"), + StatisticsType::RowCount => "row_count".to_string(), + }; + + let stat_column = phys_expr::Column::new(&stat_column_name, idx); + + // only add statistics column if not previously added + if need_to_insert { + // may be null if statistics are not present + let nullable = true; + let stat_field = + Field::new(stat_column.name(), field.data_type().clone(), nullable); + self.columns.push((column.clone(), stat_type, stat_field)); + } + rewrite_column_expr(Arc::clone(column_expr), column, &stat_column) + } + + /// rewrite col --> col_min + fn min_column_expr( + &mut self, + column: &phys_expr::Column, + column_expr: &Arc<dyn PhysicalExpr>, + field: &Field, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.stat_column_expr(column, column_expr, field, StatisticsType::Min) + } + + /// rewrite col --> col_max + fn max_column_expr( + &mut self, + column: &phys_expr::Column, + column_expr: &Arc<dyn PhysicalExpr>, + field: &Field, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.stat_column_expr(column, column_expr, field, StatisticsType::Max) + } + + /// rewrite col --> col_null_count + fn null_count_column_expr( + &mut self, + column: &phys_expr::Column, + column_expr: &Arc<dyn PhysicalExpr>, + field: &Field, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.stat_column_expr(column, column_expr, field, StatisticsType::NullCount) + } + + /// rewrite col --> col_row_count + fn row_count_column_expr( + &mut self, + column: &phys_expr::Column, + column_expr: &Arc<dyn PhysicalExpr>, + field: &Field, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.stat_column_expr(column, column_expr, field, StatisticsType::RowCount) + } +} + +impl From<Vec<(phys_expr::Column, StatisticsType, Field)>> for RequiredColumns { + fn from(columns: Vec<(phys_expr::Column, StatisticsType, Field)>) -> Self { + Self { columns } + } +} + +/// Build a RecordBatch from a list of statistics, creating arrays, +/// with one row for each PruningStatistics and columns specified in +/// the required_columns parameter. +/// +/// For example, if the requested columns are +/// ```text +/// ("s1", Min, Field:s1_min) +/// ("s2", Max, field:s2_max) +/// ``` +/// +/// And the input statistics had +/// ```text +/// S1(Min: 5, Max: 10) +/// S2(Min: 99, Max: 1000) +/// S3(Min: 1, Max: 2) +/// ``` +/// +/// Then this function would build a record batch with 2 columns and +/// one row s1_min and s2_max as follows (s3 is not requested): +/// +/// ```text +/// s1_min | s2_max +/// -------+-------- +/// 5 | 1000 +/// ``` +fn build_statistics_record_batch<S: PruningStatistics + ?Sized>( + statistics: &S, + required_columns: &RequiredColumns, +) -> Result<RecordBatch> { + let mut arrays = Vec::<ArrayRef>::new(); + // For each needed statistics column: + for (column, statistics_type, stat_field) in required_columns.iter() { + let column = Column::from_name(column.name()); + let data_type = stat_field.data_type(); + + let num_containers = statistics.num_containers(); + + let array = match statistics_type { + StatisticsType::Min => statistics.min_values(&column), + StatisticsType::Max => statistics.max_values(&column), + StatisticsType::NullCount => statistics.null_counts(&column), + StatisticsType::RowCount => statistics.row_counts(), + }; + let array = array.unwrap_or_else(|| new_null_array(data_type, num_containers)); + + assert_eq_or_internal_err!( + num_containers, + array.len(), + "mismatched statistics length. Expected {}, got {}", + num_containers, + array.len() + ); + + // cast statistics array to required data type (e.g. parquet + // provides timestamp statistics as "Int64") + let array = arrow::compute::cast(&array, data_type)?; + + arrays.push(array); + } + + let schema = Arc::new(required_columns.schema()); + // provide the count in case there were no needed statistics + let mut options = RecordBatchOptions::default(); + options.row_count = Some(statistics.num_containers()); + + trace!("Creating statistics batch for {required_columns:#?} with {arrays:#?}"); + + RecordBatch::try_new_with_options(schema, arrays, &options).map_err(|err| { + plan_datafusion_err!("Can not create statistics record batch: {err}") + }) +} + +struct PruningExpressionBuilder<'a> { + column: phys_expr::Column, + column_expr: Arc<dyn PhysicalExpr>, + op: Operator, + scalar_expr: Arc<dyn PhysicalExpr>, + field: &'a Field, + required_columns: &'a mut RequiredColumns, +} + +impl<'a> PruningExpressionBuilder<'a> { + fn try_new( + left: &'a Arc<dyn PhysicalExpr>, + right: &'a Arc<dyn PhysicalExpr>, + left_columns: ColumnReferenceCount, + right_columns: ColumnReferenceCount, + op: Operator, + schema: &'a SchemaRef, + required_columns: &'a mut RequiredColumns, + ) -> Result<Self> { + // find column name; input could be a more complicated expression + let (column_expr, scalar_expr, column, correct_operator) = match ( + left_columns, + right_columns, + ) { + (ColumnReferenceCount::One(column), ColumnReferenceCount::Zero) => { + (left, right, column, op) + } + (ColumnReferenceCount::Zero, ColumnReferenceCount::One(column)) => { + (right, left, column, reverse_operator(op)?) + } + (ColumnReferenceCount::One(_), ColumnReferenceCount::One(_)) => { + // both sides have one column - not supported + return plan_err!( + "Expression not supported for pruning: left has 1 column, right has 1 column" + ); + } + (ColumnReferenceCount::Zero, ColumnReferenceCount::Zero) => { + // both sides are literals - should be handled before calling try_new + return plan_err!( + "Pruning literal expressions is not supported, please call PhysicalExprSimplifier first" + ); + } + (ColumnReferenceCount::Many, _) | (_, ColumnReferenceCount::Many) => { + return plan_err!( + "Expression not supported for pruning: left or right has multiple columns" + ); + } + }; + + let df_schema = DFSchema::try_from(Arc::clone(schema))?; + let (column_expr, correct_operator, scalar_expr) = rewrite_expr_to_prunable( + column_expr, + correct_operator, + scalar_expr, + df_schema, + )?; + let field = match schema.column_with_name(column.name()) { + Some((_, f)) => f, + _ => { + return plan_err!("Field not found in schema"); + } + }; + + Ok(Self { + column, + column_expr, + op: correct_operator, + scalar_expr, + field, + required_columns, + }) + } + + fn op(&self) -> Operator { + self.op + } + + fn scalar_expr(&self) -> &Arc<dyn PhysicalExpr> { + &self.scalar_expr + } + + fn min_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> { + self.required_columns + .min_column_expr(&self.column, &self.column_expr, self.field) + } + + fn max_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> { + self.required_columns + .max_column_expr(&self.column, &self.column_expr, self.field) + } + + /// This function is to simply retune the `null_count` physical expression no matter what the + /// predicate expression is + /// + /// i.e., x > 5 => x_null_count, + /// cast(x as int) < 10 => x_null_count, + /// try_cast(x as float) < 10.0 => x_null_count + fn null_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> { + // Retune to [`phys_expr::Column`] + let column_expr = Arc::new(self.column.clone()) as _; + + // null_count is DataType::UInt64, which is different from the column's data type (i.e. self.field) + let null_count_field = &Field::new(self.field.name(), DataType::UInt64, true); + + self.required_columns.null_count_column_expr( + &self.column, + &column_expr, + null_count_field, + ) + } + + /// This function is to simply retune the `row_count` physical expression no matter what the + /// predicate expression is + /// + /// i.e., x > 5 => x_row_count, + /// cast(x as int) < 10 => x_row_count, + /// try_cast(x as float) < 10.0 => x_row_count + fn row_count_column_expr(&mut self) -> Result<Arc<dyn PhysicalExpr>> { + // Retune to [`phys_expr::Column`] + let column_expr = Arc::new(self.column.clone()) as _; + + // row_count is DataType::UInt64, which is different from the column's data type (i.e. self.field) + let row_count_field = &Field::new(self.field.name(), DataType::UInt64, true); + + self.required_columns.row_count_column_expr( + &self.column, + &column_expr, + row_count_field, + ) + } +} + +/// This function is designed to rewrite the column_expr to +/// ensure the column_expr is monotonically increasing. +/// +/// For example, +/// 1. `col > 10` +/// 2. `-col > 10` should be rewritten to `col < -10` +/// 3. `!col = true` would be rewritten to `col = !true` +/// 4. `abs(a - 10) > 0` not supported +/// 5. `cast(can_prunable_expr) > 10` +/// 6. `try_cast(can_prunable_expr) > 10` +/// +/// More rewrite rules are still in progress. +fn rewrite_expr_to_prunable( + column_expr: &PhysicalExprRef, + op: Operator, + scalar_expr: &PhysicalExprRef, + schema: DFSchema, +) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> { + if !is_compare_op(op) { + return plan_err!("rewrite_expr_to_prunable only support compare expression"); + } + + if column_expr.downcast_ref::<phys_expr::Column>().is_some() { + // `col op lit()` + Ok((Arc::clone(column_expr), op, Arc::clone(scalar_expr))) + } else if let Some(cast) = column_expr.downcast_ref::<phys_expr::CastExpr>() { + // `cast(col) op lit()` + let (left, op, right) = rewrite_cast_child_to_prunable( + cast.expr(), + cast.cast_type(), + op, + scalar_expr, + schema, + )?; + let left = Arc::new(phys_expr::CastExpr::new_with_target_field( + left, + Arc::clone(cast.target_field()), + None, + )); + // PruningPredicate does not support pruning on nested fields yet. + // End-to-end nested-field pruning also requires Parquet statistics + // extraction to agree with PruningPredicate on a stats representation + // for nested field expressions. + Ok((left, op, right)) + } else if let Some(try_cast) = column_expr.downcast_ref::<phys_expr::TryCastExpr>() { + // `try_cast(col) op lit()` + let (left, op, right) = rewrite_cast_child_to_prunable( + try_cast.expr(), + try_cast.cast_type(), + op, + scalar_expr, + schema, + )?; + let left = Arc::new(phys_expr::TryCastExpr::new( + left, + try_cast.cast_type().clone(), + )); + Ok((left, op, right)) + } else if let Some(neg) = column_expr.downcast_ref::<phys_expr::NegativeExpr>() { + // `-col > lit()` --> `col < -lit()` + let (left, op, right) = + rewrite_expr_to_prunable(neg.arg(), op, scalar_expr, schema)?; + let right = Arc::new(phys_expr::NegativeExpr::new(right)); + Ok((left, reverse_operator(op)?, right)) + } else if let Some(not) = column_expr.downcast_ref::<phys_expr::NotExpr>() { + // `!col = true` --> `col = !true` + if !matches!( + op, + Operator::Eq + | Operator::NotEq + | Operator::IsDistinctFrom + | Operator::IsNotDistinctFrom + ) { + return plan_err!( + "Not with operator other than Eq / NotEq / IsDistinctFrom / IsNotDistinctFrom is not supported" + ); + } + if not.arg().downcast_ref::<phys_expr::Column>().is_some() { + let left = Arc::clone(not.arg()); + let right = Arc::new(phys_expr::NotExpr::new(Arc::clone(scalar_expr))); + Ok((left, reverse_operator(op)?, right)) + } else { + plan_err!("Not with complex expression {column_expr:?} is not supported") + } + } else { + plan_err!("column expression {column_expr:?} is not supported") + } +} + +fn rewrite_cast_child_to_prunable( + cast_child_expr: &PhysicalExprRef, + cast_type: &DataType, + op: Operator, + scalar_expr: &PhysicalExprRef, + schema: DFSchema, +) -> Result<(PhysicalExprRef, Operator, PhysicalExprRef)> { + verify_support_type_for_prune( + &cast_child_expr.data_type(schema.as_arrow())?, + cast_type, + )?; + rewrite_expr_to_prunable(cast_child_expr, op, scalar_expr, schema) +} + +fn is_compare_op(op: Operator) -> bool { + matches!( + op, + Operator::Eq + | Operator::NotEq + | Operator::Lt + | Operator::LtEq + | Operator::Gt + | Operator::GtEq + | Operator::IsDistinctFrom + | Operator::IsNotDistinctFrom + | Operator::LikeMatch + | Operator::NotLikeMatch + ) +} + +// The pruning logic is based on the comparing the min/max bounds. +// Must make sure the two type has order. +// For example, casts from string to numbers is not correct. +// Because the "13" is less than "3" with UTF8 comparison order. +fn verify_support_type_for_prune(from_type: &DataType, to_type: &DataType) -> Result<()> { + // Dictionary casts are always supported as long as the value types are supported + let from_type = match from_type { + DataType::Dictionary(_, t) => { + return verify_support_type_for_prune(t.as_ref(), to_type); + } + _ => from_type, + }; + let to_type = match to_type { + DataType::Dictionary(_, t) => { + return verify_support_type_for_prune(from_type, t.as_ref()); + } + _ => to_type, + }; + // If both types are strings or both are not strings (number, timestamp, etc) + // then we can compare them. + // PruningPredicate does not support casting of strings to numbers and such. + if from_type.is_string() == to_type.is_string() { + Ok(()) + } else { + plan_err!( + "Try Cast/Cast with from type {from_type} to type {to_type} is not supported" + ) + } +} + +/// replaces a column with an old name with a new name in an expression +fn rewrite_column_expr( + e: Arc<dyn PhysicalExpr>, + column_old: &phys_expr::Column, + column_new: &phys_expr::Column, +) -> Result<Arc<dyn PhysicalExpr>> { + e.transform(|expr| { + if let Some(column) = expr.downcast_ref::<phys_expr::Column>() + && column == column_old + { + return Ok(Transformed::yes(Arc::new(column_new.clone()))); + } + + Ok(Transformed::no(expr)) + }) + .data() +} + +fn reverse_operator(op: Operator) -> Result<Operator> { + op.swap().ok_or_else(|| { + internal_datafusion_err!( + "Could not reverse operator {op} while building pruning predicate" + ) + }) +} + +/// Given a column reference to `column`, returns a pruning +/// expression in terms of the min and max that will evaluate to true +/// if the column may contain values, and false if definitely does not +/// contain values +fn build_single_column_expr( + column: &phys_expr::Column, + schema: &Schema, + required_columns: &mut RequiredColumns, + is_not: bool, // if true, treat as !col +) -> Option<Arc<dyn PhysicalExpr>> { + let field = schema.field_with_name(column.name()).ok()?; + + if *field.data_type() == DataType::Boolean { + let col_ref = Arc::new(column.clone()) as _; + + let min = required_columns + .min_column_expr(column, &col_ref, field) + .ok()?; + let max = required_columns + .max_column_expr(column, &col_ref, field) + .ok()?; + + // remember -- we want an expression that is: + // TRUE: if there may be rows that match + // FALSE: if there are no rows that match + if is_not { + // The only way we know a column couldn't match is if both the min and max are true + // !(min && max) + Some(Arc::new(phys_expr::NotExpr::new(Arc::new( + phys_expr::BinaryExpr::new(min, Operator::And, max), + )))) + } else { + // the only way we know a column couldn't match is if both the min and max are false + // !(!min && !max) --> min || max + Some(Arc::new(phys_expr::BinaryExpr::new(min, Operator::Or, max))) + } + } else { + None + } +} + +/// Given an expression reference to `expr`, if `expr` is a column expression, +/// returns a pruning expression in terms of IsNull that will evaluate to true +/// if the column may contain null, and false if definitely does not +/// contain null. +/// If `with_not` is true, build a pruning expression for `col IS NOT NULL`: `col_count != col_null_count` +/// The pruning expression evaluates to true ONLY if the column definitely CONTAINS +/// at least one NULL value. In this case we can know that `IS NOT NULL` can not be true and +/// thus can prune the row group / value +fn build_is_null_column_expr( + expr: &Arc<dyn PhysicalExpr>, + schema: &Schema, + required_columns: &mut RequiredColumns, + with_not: bool, +) -> Option<Arc<dyn PhysicalExpr>> { + if let Some(col) = expr.downcast_ref::<phys_expr::Column>() { + let field = schema.field_with_name(col.name()).ok()?; + + let null_count_field = &Field::new(field.name(), DataType::UInt64, true); + if with_not { + if let Ok(row_count_expr) = + required_columns.row_count_column_expr(col, expr, null_count_field) + { + required_columns + .null_count_column_expr(col, expr, null_count_field) + .map(|null_count_column_expr| { + // IsNotNull(column) => null_count != row_count + Arc::new(phys_expr::BinaryExpr::new( + null_count_column_expr, + Operator::NotEq, + row_count_expr, + )) as _ + }) + .ok() + } else { + None + } + } else { + required_columns + .null_count_column_expr(col, expr, null_count_field) + .map(|null_count_column_expr| { + // IsNull(column) => null_count > 0 + Arc::new(phys_expr::BinaryExpr::new( + null_count_column_expr, + Operator::Gt, + Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))), + )) as _ + }) + .ok() + } + } else { + None + } +} + +/// Default maximum number of entries in an `IN (...)` list that will be +/// rewritten into a chain of per-value min/max checks by +/// `build_predicate_expression`. Callers threading a [`PredicateRewriter`] +/// can override this via [`PredicateRewriter::with_max_in_list_size`], and +/// query engines can wire it from the +/// `datafusion.execution.parquet.max_in_list_size` config option. +pub const MAX_IN_LIST_SIZE: usize = 20; + +/// Statistics expressions for `column`: its min, its max, and "not every row is +/// NULL". Appends to `required_columns`, rolling the appends back if any part +/// cannot be built, since [`RequiredColumns::stat_column_expr`] only appends. +fn min_max_non_null_exprs( + column: &phys_expr::Column, + column_expr: &Arc<dyn PhysicalExpr>, + field: &Field, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option<(PhysicalExprRef, PhysicalExprRef, PhysicalExprRef)> { + let appended = required_columns.columns.len(); + let statistics = (|| { + let min = required_columns + .min_column_expr(column, column_expr, field) + .ok()?; + let max = required_columns + .max_column_expr(column, column_expr, field) + .ok()?; + // A cast on `column_expr` (widening it towards min/max's target type) + // does not affect nullability, so check it on the bare column beneath. + let bare_expr = match column_expr.downcast_ref::<phys_expr::CastExpr>() { + Some(cast) => cast.expr(), + None => column_expr, + }; + let non_null = + build_is_null_column_expr(bare_expr, schema, required_columns, true)?; + Some((min, max, non_null)) + })(); + if statistics.is_none() { + required_columns.columns.truncate(appended); + } + statistics +} + +/// Builds a "may match" expression for a [`HashTableLookupExpr`] (a large join +/// build side pushed down as an opaque hash-table lookup): tests container +/// min/max stats against the build side's key-range bitmap, which is fixed size +/// and built once with the join's hash table. +fn build_hash_lookup_pruning_expr( + lookup: &HashTableLookupExpr, + schema: &Schema, + required_columns: &mut RequiredColumns, +) -> Option<Arc<dyn PhysicalExpr>> { + let bitmap = Arc::clone(lookup.pruning_bitmap()?); + let on_columns = lookup.children(); + let [column_expr] = on_columns[..] else { + return None; + }; + // The schema adapter often wraps the column in a same-type or widening cast + // (a nullable table column over a REQUIRED file column, or a narrower file + // type). `stat_column_expr` rewrites only the inner `Column`, keeping it. + let column = match column_expr.downcast_ref::<phys_expr::CastExpr>() { + Some(cast) => { + let column = cast.expr().downcast_ref::<phys_expr::Column>()?; + let from = schema.fields().get(column.index())?.data_type(); + // 55.1.0's widening-cast helper predates unsigned-to-larger-signed + // support. These additional integer casts preserve every value, + // ordering, and null; keep other unrecognized casts conservative. + let widens_unsigned = matches!( + (from, cast.cast_type()), + ( + DataType::UInt8, + DataType::Int16 | DataType::Int32 | DataType::Int64 + ) | (DataType::UInt16, DataType::Int32 | DataType::Int64) + | (DataType::UInt32, DataType::Int64) + ); + (cast.is_bigger_cast(from) || widens_unsigned).then_some(column)? + } + None => column_expr.downcast_ref::<phys_expr::Column>()?, + }; + let field = schema.fields().get(column.index())?; + if field.name() != column.name() { + return None; + } + let (min, max, non_null) = + min_max_non_null_exprs(column, column_expr, field, schema, required_columns)?; + Some(Arc::new(phys_expr::BinaryExpr::new( + non_null, + Operator::And, + Arc::new(KeyRangeBitmapPruningExpr { min, max, bitmap }), + ))) +} + +/// Rewrite a predicate expression in terms of statistics (min/max/null_counts) +/// for use as a [`PruningPredicate`]. +pub struct PredicateRewriter { + unhandled_hook: Arc<dyn UnhandledPredicateHook>, + max_in_list_size: usize, +} + +impl Default for PredicateRewriter { + fn default() -> Self { + Self { + unhandled_hook: Arc::new(ConstantUnhandledPredicateHook::default()), + max_in_list_size: MAX_IN_LIST_SIZE, + } + } +} + +impl PredicateRewriter { + /// Create a new `PredicateRewriter` + pub fn new() -> Self { + Self::default() + } + + /// Set the unhandled hook to be used when a predicate can not be rewritten + pub fn with_unhandled_hook( + mut self, + unhandled_hook: Arc<dyn UnhandledPredicateHook>, + ) -> Self { + self.unhandled_hook = unhandled_hook; + self + } + + /// Set the maximum size of an `IN (...)` list that will be rewritten into a + /// chain of per-value statistics checks. Lists longer than this fall back + /// to the unhandled-predicate hook (typically "keep the container"), + /// effectively skipping container-level pruning for large IN lists. + /// + /// The default (see [`MAX_IN_LIST_SIZE`]) preserves the + /// historical behaviour. Callers wiring config through can override via + /// `datafusion.execution.max_in_list_size`. + pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self { + self.max_in_list_size = max_in_list_size; + self + } + + /// Translate logical filter expression into pruning predicate + /// expression that will evaluate to FALSE if it can be determined no + /// rows between the min/max values could pass the predicates. + /// + /// Any predicates that can not be translated will be passed to `unhandled_hook`. + /// + /// Returns the pruning predicate as an [`PhysicalExpr`] + /// + /// Notice: `IN (...)` lists longer than `max_in_list_size` (default + /// [`MAX_IN_LIST_SIZE`]) fall back to calling `unhandled_hook`. + pub fn rewrite_predicate_to_statistics_predicate( + &self, + expr: &Arc<dyn PhysicalExpr>, + schema: &Schema, + ) -> Arc<dyn PhysicalExpr> { + let mut required_columns = RequiredColumns::new(); + build_predicate_expression( + expr, + &Arc::new(schema.clone()), + &mut required_columns, + &self.unhandled_hook, + self.max_in_list_size, + ) + } +} + +/// Translate logical filter expression into pruning predicate +/// expression that will evaluate to FALSE if it can be determined no +/// rows between the min/max values could pass the predicates. +/// +/// Any predicates that can not be translated will be passed to `unhandled_hook`. +/// +/// Returns the pruning predicate as an [`PhysicalExpr`] +/// +/// `max_in_list_size` is the largest `IN (...)` list that will be rewritten +/// into a chain of per-value statistics checks; longer lists fall back to +/// `unhandled_hook`. +fn build_predicate_expression( + expr: &Arc<dyn PhysicalExpr>, + schema: &SchemaRef, + required_columns: &mut RequiredColumns, + unhandled_hook: &Arc<dyn UnhandledPredicateHook>, + max_in_list_size: usize, +) -> Arc<dyn PhysicalExpr> { + if is_always_false(expr) { + // Shouldn't return `unhandled_hook.handle(expr)` + // Because it will transfer false to true. + return Arc::clone(expr); + } + // predicate expression can only be a binary expression + if let Some(is_null) = expr.downcast_ref::<phys_expr::IsNullExpr>() { + return build_is_null_column_expr(is_null.arg(), schema, required_columns, false) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + if let Some(is_not_null) = expr.downcast_ref::<phys_expr::IsNotNullExpr>() { + return build_is_null_column_expr( + is_not_null.arg(), + schema, + required_columns, + true, + ) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + if let Some(col) = expr.downcast_ref::<phys_expr::Column>() { + return build_single_column_expr(col, schema, required_columns, false) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + if let Some(not) = expr.downcast_ref::<phys_expr::NotExpr>() { + // match !col (don't do so recursively) + if let Some(col) = not.arg().downcast_ref::<phys_expr::Column>() { + return build_single_column_expr(col, schema, required_columns, true) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } else { + return unhandled_hook.handle(expr); + } + } + if let Some(in_list) = expr.downcast_ref::<phys_expr::InListExpr>() { + if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size { + let eq_op = if in_list.negated() { + Operator::NotEq + } else { + Operator::Eq + }; + let re_op = if in_list.negated() { + Operator::And + } else { + Operator::Or + }; + let change_expr = in_list + .list() + .iter() + .map(|e| { + Arc::new(phys_expr::BinaryExpr::new( + Arc::clone(in_list.expr()), + eq_op, + Arc::clone(e), + )) as _ + }) + .reduce(|a, b| Arc::new(phys_expr::BinaryExpr::new(a, re_op, b)) as _) + .unwrap(); + return build_predicate_expression( + &change_expr, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); + } else { + return unhandled_hook.handle(expr); + } + } + if let Some(lookup) = expr.downcast_ref::<HashTableLookupExpr>() { + return build_hash_lookup_pruning_expr(lookup, schema, required_columns) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + // A partitioned hash join hides its per-partition filters under a `CASE` on the + // repartition hash. A row takes exactly one branch, so a container may match + // only if some branch may: the branches' disjunction is a sound relaxation, and + // the `WHEN`s (a hash, which no statistics describe) can be dropped. + if let Some(case) = expr.downcast_ref::<phys_expr::CaseExpr>() { + // Only a Boolean `CASE` is a predicate; anything else is a value for + // whatever compares it to handle. + if !matches!(case.data_type(schema), Ok(DataType::Boolean)) { + return unhandled_hook.handle(expr); + } + // Without an `ELSE`, unmatched rows are UNKNOWN: not a match, but not + // FALSE either, which full-match inference must be able to tell apart. + if case.else_expr().is_none() { + return unhandled_hook.handle(expr); + } + return case + .when_then_expr() + .iter() + .map(|(_, then)| then) + .chain(case.else_expr()) + .map(|branch| { + build_predicate_expression( + branch, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ) + }) + .reduce(|acc, branch| { + Arc::new(phys_expr::BinaryExpr::new(acc, Operator::Or, branch)) as _ + }) + .unwrap_or_else(|| unhandled_hook.handle(expr)); + } + + let (left, op, right) = { + if let Some(bin_expr) = expr.downcast_ref::<phys_expr::BinaryExpr>() { + ( + Arc::clone(bin_expr.left()), + *bin_expr.op(), + Arc::clone(bin_expr.right()), + ) + } else if let Some(like_expr) = expr.downcast_ref::<phys_expr::LikeExpr>() { + if like_expr.case_insensitive() { + return unhandled_hook.handle(expr); + } + let op = match (like_expr.negated(), like_expr.case_insensitive()) { + (false, false) => Operator::LikeMatch, + (true, false) => Operator::NotLikeMatch, + (false, true) => Operator::ILikeMatch, + (true, true) => Operator::NotILikeMatch, + }; + ( + Arc::clone(like_expr.expr()), + op, + Arc::clone(like_expr.pattern()), + ) + } else { + return unhandled_hook.handle(expr); + } + }; + + if op == Operator::And || op == Operator::Or { + let left_expr = build_predicate_expression( + &left, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); + let right_expr = build_predicate_expression( + &right, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); + // simplify boolean expression if applicable + let expr = match (&left_expr, op, &right_expr) { + (left, Operator::And, right) + if is_always_false(left) || is_always_false(right) => + { + Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(false)))) + } + (left, Operator::And, _) if is_always_true(left) => right_expr, + (_, Operator::And, right) if is_always_true(right) => left_expr, + (left, Operator::Or, right) + if is_always_true(left) || is_always_true(right) => + { + Arc::new(phys_expr::Literal::new(ScalarValue::Boolean(Some(true)))) + } + (left, Operator::Or, _) if is_always_false(left) => right_expr, + (_, Operator::Or, right) if is_always_false(right) => left_expr, + + _ => Arc::new(phys_expr::BinaryExpr::new(left_expr, op, right_expr)), + }; + return expr; + } + + let left_columns = ColumnReferenceCount::from_expression(&left); + let right_columns = ColumnReferenceCount::from_expression(&right); + let expr_builder = PruningExpressionBuilder::try_new( + &left, + &right, + left_columns, + right_columns, + op, + schema, + required_columns, + ); + let mut expr_builder = match expr_builder { + Ok(builder) => builder, + // allow partial failure in predicate expression generation + // this can still produce a useful predicate when multiple conditions are joined using AND + Err(e) => { + debug!("Error building pruning expression: {e}"); + return unhandled_hook.handle(expr); + } + }; + + build_statistics_expr(&mut expr_builder) + .unwrap_or_else(|_| unhandled_hook.handle(expr)) +} + +/// Count of distinct column references in an expression. +/// This is the same as [`collect_columns`] but optimized to stop counting +/// once more than one distinct column is found. +/// +/// For example, in expression `col1 + col2`, the count is `Many`. +/// In expression `col1 + 5`, the count is `One`. +/// In expression `5 + 10`, the count is `Zero`. +/// +/// [`collect_columns`]: datafusion_physical_expr::utils::collect_columns +#[derive(Debug, PartialEq, Eq)] +enum ColumnReferenceCount { + /// no column references + Zero, + /// Only one column reference + One(phys_expr::Column), + /// More than one column reference + Many, +} + +impl ColumnReferenceCount { + /// Count the number of distinct column references in an expression + fn from_expression(expr: &Arc<dyn PhysicalExpr>) -> Self { + let mut seen = HashSet::<phys_expr::Column>::new(); + expr.apply(|expr| { + if let Some(column) = expr.downcast_ref::<phys_expr::Column>() { + seen.insert(column.clone()); + if seen.len() > 1 { + return Ok(TreeNodeRecursion::Stop); + } + } + Ok(TreeNodeRecursion::Continue) + }) + // pre_visit always returns OK, so this will always too + .expect("no way to return error during recursion"); + match seen.len() { + 0 => ColumnReferenceCount::Zero, + 1 => ColumnReferenceCount::One( + seen.into_iter().next().expect("just checked len==1"), + ), + _ => ColumnReferenceCount::Many, + } + } +} + +fn build_statistics_expr( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + let statistics_expr: Arc<dyn PhysicalExpr> = match expr_builder.op() { + Operator::NotEq => build_ne_statistics_expr(expr_builder)?, + Operator::Eq => { + // column = literal => (min, max) = literal => min <= literal && literal <= max + // (column / 2) = 4 => (column_min / 2) <= 4 && 4 <= (column_max / 2) + build_eq_statistics_expr(expr_builder)? + } + Operator::IsDistinctFrom => return build_is_distinct_from(expr_builder), + Operator::IsNotDistinctFrom => return build_is_not_distinct_from(expr_builder), + Operator::NotLikeMatch => build_not_like_match(expr_builder)?, + Operator::LikeMatch => build_like_match(expr_builder).ok_or_else(|| { + plan_datafusion_err!( + "LIKE expression with wildcard at the beginning is not supported" + ) + })?, + Operator::Gt => { + // column > literal => (min, max) > literal => max > literal + Arc::new(phys_expr::BinaryExpr::new( + expr_builder.max_column_expr()?, + Operator::Gt, + Arc::clone(expr_builder.scalar_expr()), + )) + } + Operator::GtEq => { + // column >= literal => (min, max) >= literal => max >= literal + Arc::new(phys_expr::BinaryExpr::new( + expr_builder.max_column_expr()?, + Operator::GtEq, + Arc::clone(expr_builder.scalar_expr()), + )) + } + Operator::Lt => { + // column < literal => (min, max) < literal => min < literal + Arc::new(phys_expr::BinaryExpr::new( + expr_builder.min_column_expr()?, + Operator::Lt, + Arc::clone(expr_builder.scalar_expr()), + )) + } + Operator::LtEq => { + // column <= literal => (min, max) <= literal => min <= literal + Arc::new(phys_expr::BinaryExpr::new( + expr_builder.min_column_expr()?, + Operator::LtEq, + Arc::clone(expr_builder.scalar_expr()), + )) + } + // other expressions are not supported + _ => { + return plan_err!( + "expressions other than (neq, eq, gt, gteq, lt, lteq) are not supported" + ); + } + }; + let statistics_expr = wrap_null_count_check_expr(statistics_expr, expr_builder)?; + Ok(statistics_expr) +} + +fn binary_expr( + left: Arc<dyn PhysicalExpr>, + op: Operator, + right: Arc<dyn PhysicalExpr>, +) -> Arc<dyn PhysicalExpr> { + Arc::new(phys_expr::BinaryExpr::new(left, op, right)) +} + +fn and_expr( + left: Arc<dyn PhysicalExpr>, + right: Arc<dyn PhysicalExpr>, +) -> Arc<dyn PhysicalExpr> { + binary_expr(left, Operator::And, right) +} + +fn or_expr( + left: Arc<dyn PhysicalExpr>, + right: Arc<dyn PhysicalExpr>, +) -> Arc<dyn PhysicalExpr> { + binary_expr(left, Operator::Or, right) +} + +fn build_eq_statistics_expr( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + let min_column_expr = expr_builder.min_column_expr()?; + let max_column_expr = expr_builder.max_column_expr()?; + Ok(and_expr( + binary_expr( + min_column_expr, + Operator::LtEq, + Arc::clone(expr_builder.scalar_expr()), + ), + binary_expr( + Arc::clone(expr_builder.scalar_expr()), + Operator::LtEq, + max_column_expr, + ), + )) +} + +fn build_ne_statistics_expr( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + let min_column_expr = expr_builder.min_column_expr()?; + let max_column_expr = expr_builder.max_column_expr()?; + Ok(or_expr( + binary_expr( + min_column_expr, + Operator::NotEq, + Arc::clone(expr_builder.scalar_expr()), + ), + binary_expr( + Arc::clone(expr_builder.scalar_expr()), + Operator::NotEq, + max_column_expr, + ), + )) +} + +fn column_has_nulls_expr( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + Ok(binary_expr( + expr_builder.null_count_column_expr()?, + Operator::Gt, + Arc::new(phys_expr::Literal::new(ScalarValue::UInt64(Some(0)))), + )) +} + +fn column_has_non_nulls_expr( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + Ok(binary_expr( + expr_builder.null_count_column_expr()?, + Operator::NotEq, + expr_builder.row_count_column_expr()?, + )) +} + +fn build_is_distinct_from( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + let scalar_expr = Arc::clone(expr_builder.scalar_expr()); + + Ok(or_expr( + and_expr( + Arc::new(phys_expr::IsNullExpr::new(Arc::clone(&scalar_expr))), + column_has_non_nulls_expr(expr_builder)?, + ), + and_expr( + Arc::new(phys_expr::IsNotNullExpr::new(scalar_expr)), + or_expr( + column_has_nulls_expr(expr_builder)?, + build_ne_statistics_expr(expr_builder)?, + ), + ), + )) +} + +fn build_is_not_distinct_from( + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + let scalar_expr = Arc::clone(expr_builder.scalar_expr()); + + Ok(or_expr( + and_expr( + Arc::new(phys_expr::IsNullExpr::new(Arc::clone(&scalar_expr))), + column_has_nulls_expr(expr_builder)?, + ), + and_expr( + Arc::new(phys_expr::IsNotNullExpr::new(scalar_expr)), + and_expr( + column_has_non_nulls_expr(expr_builder)?, + build_eq_statistics_expr(expr_builder)?, + ), + ), + )) +} + +/// returns the string literal of the scalar value if it is a string +fn unpack_string(s: &ScalarValue) -> Option<&str> { + s.try_as_str().flatten() +} + +fn extract_string_literal(expr: &Arc<dyn PhysicalExpr>) -> Option<&str> { + if let Some(lit) = expr.downcast_ref::<phys_expr::Literal>() { + let s = unpack_string(lit.value())?; + return Some(s); + } + None +} + +/// Wrap a string in a `Literal` whose `ScalarValue` matches `target_type` +fn string_literal_as(value: String, target_type: &DataType) -> Arc<dyn PhysicalExpr> { + let utf8 = ScalarValue::Utf8(Some(value)); + let scalar = try_cast_literal_to_type(&utf8, target_type).unwrap_or(utf8); + Arc::new(phys_expr::Literal::new(scalar)) +} + +/// Convert `column LIKE literal` where P is a constant prefix of the literal +/// to a range check on the column: `P <= column && column < P'`, where P' is the +/// lowest string after all P* strings. +fn build_like_match( + expr_builder: &mut PruningExpressionBuilder, +) -> Option<Arc<dyn PhysicalExpr>> { + // column LIKE literal => (min, max) LIKE literal split at unescaped % => min <= split literal && split literal <= max + // column LIKE 'foo%' => min <= 'foo' && 'foo' <= max + // column LIKE 'foo\_%' => min <= 'foo_' && 'foo_' <= max (the _ is escaped) + // column LIKE 'foo\%%' => min <= 'foo%' && 'foo%' <= max (the % is escaped) + // column LIKE '%foo' => min <= '' && '' <= max => true + // column LIKE '%foo%' => min <= '' && '' <= max => true + // column LIKE 'foo' => min <= 'foo' && 'foo' <= max + + // TODO Handle ILIKE perhaps by making the min lowercase and max uppercase + // this may involve building the physical expressions that call lower() and upper() + let min_column_expr = expr_builder.min_column_expr().ok()?; + let max_column_expr = expr_builder.max_column_expr().ok()?; + let scalar_expr = expr_builder.scalar_expr(); + // Synthesized bounds must match the column type (e.g. `Utf8View`). + let target_type = expr_builder.field.data_type(); + // check that the scalar is a string literal + let s = extract_string_literal(scalar_expr)?; + // ANSI SQL specifies two wildcards: % and _. % matches zero or more characters, _ matches exactly one character. + let (decoded_prefix, rest) = split_constant_prefix(s); + let has_wildcard = !rest.is_empty(); + if has_wildcard && decoded_prefix.is_empty() { + // there's no filtering we could possibly do, return None and have this be handled by the unhandled hook + return None; + } + let (lower_bound, upper_bound) = if has_wildcard { + let incremented_prefix = increment_utf8(&decoded_prefix)?; + let lower_bound_lit = string_literal_as(decoded_prefix, target_type); + let upper_bound_lit = string_literal_as(incremented_prefix, target_type); + (lower_bound_lit, upper_bound_lit) + } else { + // the like expression is a literal and can be converted into a comparison + let bound = string_literal_as(decoded_prefix, target_type); + (Arc::clone(&bound), bound) + }; + let lower_bound_expr = Arc::new(phys_expr::BinaryExpr::new( + lower_bound, + Operator::LtEq, + Arc::clone(&max_column_expr), + )); + let upper_bound_expr = Arc::new(phys_expr::BinaryExpr::new( + Arc::clone(&min_column_expr), + Operator::LtEq, + upper_bound, + )); + let combined = Arc::new(phys_expr::BinaryExpr::new( + upper_bound_expr, + Operator::And, + lower_bound_expr, + )); + Some(combined) +} + +// For predicate `col NOT LIKE 'const_prefix%'`, we rewrite it as `(col_min NOT LIKE 'const_prefix%' OR col_max NOT LIKE 'const_prefix%')`. +// +// The intuition is that if both `col_min` and `col_max` begin with `const_prefix` that means +// **all** data in this row group begins with `const_prefix` as well (and therefore the predicate +// looking for rows that don't begin with `const_prefix` can never be true) +fn build_not_like_match( + expr_builder: &mut PruningExpressionBuilder<'_>, +) -> Result<Arc<dyn PhysicalExpr>> { + // col NOT LIKE 'const_prefix%' -> !(col_min LIKE 'const_prefix%' && col_max LIKE 'const_prefix%') -> (col_min NOT LIKE 'const_prefix%' || col_max NOT LIKE 'const_prefix%') + + let min_column_expr = expr_builder.min_column_expr()?; + let max_column_expr = expr_builder.max_column_expr()?; + + let scalar_expr = expr_builder.scalar_expr(); + + let pattern = extract_string_literal(scalar_expr).ok_or_else(|| { + plan_datafusion_err!("cannot extract literal from NOT LIKE expression") + })?; + + let (const_prefix, remaining) = split_constant_prefix(pattern); + if const_prefix.is_empty() || remaining != "%" { + // we can not handle `%` at the beginning or in the middle of the pattern + // Example: For pattern "foo%bar", the row group might include values like + // ["foobar", "food", "foodbar"], making it unsafe to prune. + // Even if the min/max values in the group (e.g., "foobar" and "foodbar") + // match the pattern, intermediate values like "food" may not + // match the full pattern "foo%bar", making pruning unsafe. + // (truncate foo%bar to foo% have same problem) + + // we can not handle pattern containing `_` + // Example: For pattern "foo_", row groups might contain ["fooa", "fooaa", "foob"], + // which means not every row is guaranteed to match the pattern. + return Err(plan_datafusion_err!( + "NOT LIKE expressions only support constant_prefix+wildcard`%`" + )); + } + + let min_col_not_like_epxr = Arc::new(phys_expr::LikeExpr::new( + true, + false, + Arc::clone(&min_column_expr), + Arc::clone(scalar_expr), + )); + + let max_col_not_like_expr = Arc::new(phys_expr::LikeExpr::new( + true, + false, + Arc::clone(&max_column_expr), + Arc::clone(scalar_expr), + )); + + Ok(Arc::new(phys_expr::BinaryExpr::new( + min_col_not_like_epxr, + Operator::Or, + max_col_not_like_expr, + ))) +} + +/// Returns unescaped constant prefix of a LIKE pattern (possibly empty) and the remaining pattern (possibly empty) +fn split_constant_prefix(pattern: &str) -> (String, &str) { + let mut prefix = String::with_capacity(pattern.len()); + let mut iter = pattern.char_indices(); + while let Some((idx, c)) = iter.next() { + match c { + '%' | '_' => return (prefix, &pattern[idx..]), + '\\' => match iter.next() { + Some((_, escaped)) => prefix.push(escaped), + None => prefix.push('\\'), + }, + _ => prefix.push(c), + } + } + (prefix, "") +} + +/// Increment a UTF8 string by one, returning `None` if it can't be incremented. +/// This makes it so that the returned string will always compare greater than the input string +/// or any other string with the same prefix. +/// This is necessary since the statistics may have been truncated: if we have a min statistic +/// of "fo" that may have originally been "foz" or anything else with the prefix "fo". +/// E.g. `increment_utf8("foo") >= "foo"` and `increment_utf8("foo") >= "fooz"` +/// In this example `increment_utf8("foo") == "fop" +fn increment_utf8(data: &str) -> Option<String> { + // Helper function to check if a character is valid to use + fn is_valid_unicode(c: char) -> bool { + let cp = c as u32; + + // Filter out non-characters (https://www.unicode.org/versions/corrigendum9.html) + if [0xFFFE, 0xFFFF].contains(&cp) || (0xFDD0..=0xFDEF).contains(&cp) { + return false; + } + + // Filter out private use area + if cp >= 0x110000 { + return false; + } + + true + } + + // Convert string to vector of code points + let mut code_points: Vec<char> = data.chars().collect(); + + // Work backwards through code points + for idx in (0..code_points.len()).rev() { + let original = code_points[idx] as u32; + + // Try incrementing the code point + if let Some(next_char) = char::from_u32(original + 1) + && is_valid_unicode(next_char) + { + code_points[idx] = next_char; + // truncate the string to the current index + code_points.truncate(idx + 1); + return Some(code_points.into_iter().collect()); + } + } + + None +} + +/// Wrap the statistics expression in a check that skips the expression if the column is all nulls. +/// +/// This is important not only as an optimization but also because statistics may not be +/// accurate for columns that are all nulls. +/// For example, for an `int` column `x` with all nulls, the min/max/null_count statistics +/// might be set to 0 and evaluating `x = 0` would incorrectly include the column. +/// +/// For example: +/// +/// `x_min <= 10 AND 10 <= x_max` +/// +/// will become +/// +/// ```sql +/// x_null_count != x_row_count AND (x_min <= 10 AND 10 <= x_max) +/// ```` +/// +/// If the column is known to be all nulls, then the expression +/// `x_null_count = x_row_count` will be true, which will cause the +/// boolean expression to return false. Therefore, prune out the container. +fn wrap_null_count_check_expr( + statistics_expr: Arc<dyn PhysicalExpr>, + expr_builder: &mut PruningExpressionBuilder, +) -> Result<Arc<dyn PhysicalExpr>> { + // (x_null_count != x_row_count) AND (<statistics_expr>) + Ok(and_expr( + column_has_non_nulls_expr(expr_builder)?, + statistics_expr, + )) +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub(crate) enum StatisticsType { + Min, + Max, + NullCount, + RowCount, +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::ops::{Not, Rem}; + + use super::*; + use datafusion_common::test_util::batches_to_string; + use datafusion_expr::{and, col, lit, or}; + use datafusion_physical_expr::utils::collect_columns; + use insta::assert_snapshot; + + use arrow::array::Decimal128Array; + use arrow::{ + array::{BinaryArray, Int32Array, Int64Array, StringArray, UInt64Array}, + datatypes::TimeUnit, + }; + use datafusion_expr::expr::InList; + use datafusion_expr::{BinaryExpr, Expr, cast, is_null, try_cast}; + use datafusion_functions_nested::expr_fn::{array_has, make_array}; + use datafusion_physical_expr::expressions::{ + self as phys_expr, DynamicFilterPhysicalExpr, + }; + use datafusion_physical_expr::planner::logical2physical; + use datafusion_physical_plan::joins::join_hash_map::JoinHashMapU32; + use datafusion_physical_plan::joins::key_range_bitmap::KeyRangeBitmap; + use datafusion_physical_plan::joins::{Map, SeededRandomState}; + use itertools::Itertools; + use rstest::rstest; + + #[derive(Debug, Default)] + /// Mock statistic provider for tests + /// + /// Each row represents the statistics for a "container" (which + /// might represent an entire parquet file, or directory of files, + /// or some other collection of data for which we had statistics) + /// + /// Note All `ArrayRefs` must be the same size. + struct ContainerStats { + min: Option<ArrayRef>, + max: Option<ArrayRef>, + /// Optional values + null_counts: Option<ArrayRef>, + row_counts: Option<ArrayRef>, + /// Optional known values (e.g. mimic a bloom filter) + /// (value, contained) + /// If present, all BooleanArrays must be the same size as min/max + contained: Vec<(HashSet<ScalarValue>, BooleanArray)>, + } + + impl ContainerStats { + fn new() -> Self { + Default::default() + } + fn new_decimal128( + min: impl IntoIterator<Item = Option<i128>>, + max: impl IntoIterator<Item = Option<i128>>, + precision: u8, + scale: i8, + ) -> Self { + Self::new() + .with_min(Arc::new( + min.into_iter() + .collect::<Decimal128Array>() + .with_precision_and_scale(precision, scale) + .unwrap(), + )) + .with_max(Arc::new( + max.into_iter() + .collect::<Decimal128Array>() + .with_precision_and_scale(precision, scale) + .unwrap(), + )) + } + + fn new_i64( + min: impl IntoIterator<Item = Option<i64>>, + max: impl IntoIterator<Item = Option<i64>>, + ) -> Self { + Self::new() + .with_min(Arc::new(min.into_iter().collect::<Int64Array>())) + .with_max(Arc::new(max.into_iter().collect::<Int64Array>())) + } + + fn new_i32( + min: impl IntoIterator<Item = Option<i32>>, + max: impl IntoIterator<Item = Option<i32>>, + ) -> Self { + Self::new() + .with_min(Arc::new(min.into_iter().collect::<Int32Array>())) + .with_max(Arc::new(max.into_iter().collect::<Int32Array>())) + } + + fn new_utf8<'a>( + min: impl IntoIterator<Item = Option<&'a str>>, + max: impl IntoIterator<Item = Option<&'a str>>, + ) -> Self { + Self::new() + .with_min(Arc::new(min.into_iter().collect::<StringArray>())) + .with_max(Arc::new(max.into_iter().collect::<StringArray>())) + } + + fn new_bool( + min: impl IntoIterator<Item = Option<bool>>, + max: impl IntoIterator<Item = Option<bool>>, + ) -> Self { + Self::new() + .with_min(Arc::new(min.into_iter().collect::<BooleanArray>())) + .with_max(Arc::new(max.into_iter().collect::<BooleanArray>())) + } + + fn min(&self) -> Option<ArrayRef> { + self.min.clone() + } + + fn max(&self) -> Option<ArrayRef> { + self.max.clone() + } + + fn null_counts(&self) -> Option<ArrayRef> { + self.null_counts.clone() + } + + fn row_counts(&self) -> Option<ArrayRef> { + self.row_counts.clone() + } + + /// return an iterator over all arrays in this statistics + fn arrays(&self) -> Vec<ArrayRef> { + let contained_arrays = self + .contained + .iter() + .map(|(_values, contained)| Arc::new(contained.clone()) as ArrayRef); + + [ + self.min.as_ref().cloned(), + self.max.as_ref().cloned(), + self.null_counts.as_ref().cloned(), + self.row_counts.as_ref().cloned(), + ] + .into_iter() + .flatten() + .chain(contained_arrays) + .collect() + } + + /// Returns the number of containers represented by this statistics This + /// picks the length of the first array as all arrays must have the same + /// length (which is verified by `assert_invariants`). + fn len(&self) -> usize { + // pick the first non zero length + self.arrays().iter().map(|a| a.len()).next().unwrap_or(0) + } + + /// Ensure that the lengths of all arrays are consistent + fn assert_invariants(&self) { + let mut prev_len = None; + + for len in self.arrays().iter().map(|a| a.len()) { + // Get a length, if we don't already have one + match prev_len { + None => { + prev_len = Some(len); + } + Some(prev_len) => { + assert_eq!(prev_len, len); + } + } + } + } + + /// Add min values + fn with_min(mut self, min: ArrayRef) -> Self { + self.min = Some(min); + self + } + + /// Add max values + fn with_max(mut self, max: ArrayRef) -> Self { + self.max = Some(max); + self + } + + /// Add null counts. There must be the same number of null counts as + /// there are containers + fn with_null_counts( + mut self, + counts: impl IntoIterator<Item = Option<u64>>, + ) -> Self { + let null_counts: ArrayRef = + Arc::new(counts.into_iter().collect::<UInt64Array>()); + + self.assert_invariants(); + self.null_counts = Some(null_counts); + self + } + + /// Add row counts. There must be the same number of row counts as + /// there are containers + fn with_row_counts( + mut self, + counts: impl IntoIterator<Item = Option<u64>>, + ) -> Self { + let row_counts: ArrayRef = + Arc::new(counts.into_iter().collect::<UInt64Array>()); + + self.assert_invariants(); + self.row_counts = Some(row_counts); + self + } + + /// Add contained information. + #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key + pub fn with_contained( + mut self, + values: impl IntoIterator<Item = ScalarValue>, + contained: impl IntoIterator<Item = Option<bool>>, + ) -> Self { + let contained: BooleanArray = contained.into_iter().collect(); + let values: HashSet<_> = values.into_iter().collect(); + + self.contained.push((values, contained)); + self.assert_invariants(); + self + } + + /// get any contained information for the specified values + #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key + fn contained(&self, find_values: &HashSet<ScalarValue>) -> Option<BooleanArray> { + // find the one with the matching values + self.contained + .iter() + .find(|(values, _contained)| values == find_values) + .map(|(_values, contained)| contained.clone()) + } + } + + #[derive(Debug, Default)] + struct TestStatistics { + // key: column name + stats: HashMap<Column, ContainerStats>, + } + + impl TestStatistics { + fn new() -> Self { + Self::default() + } + + fn with( + mut self, + name: impl Into<String>, + container_stats: ContainerStats, + ) -> Self { + let col = Column::from_name(name.into()); + self.stats.insert(col, container_stats); + self + } + + /// Add null counts for the specified column. + /// There must be the same number of null counts as + /// there are containers + fn with_null_counts( + mut self, + name: impl Into<String>, + counts: impl IntoIterator<Item = Option<u64>>, + ) -> Self { + let col = Column::from_name(name.into()); + + // take stats out and update them + let container_stats = self + .stats + .remove(&col) + .unwrap_or_default() + .with_null_counts(counts); + + // put stats back in + self.stats.insert(col, container_stats); + self + } + + /// Add row counts for the specified column. + /// There must be the same number of row counts as + /// there are containers + fn with_row_counts( + mut self, + name: impl Into<String>, + counts: impl IntoIterator<Item = Option<u64>>, + ) -> Self { + let col = Column::from_name(name.into()); + + // take stats out and update them + let container_stats = self + .stats + .remove(&col) + .unwrap_or_default() + .with_row_counts(counts); + + // put stats back in + self.stats.insert(col, container_stats); + self + } + + /// Add contained information for the specified column. + fn with_contained( + mut self, + name: impl Into<String>, + values: impl IntoIterator<Item = ScalarValue>, + contained: impl IntoIterator<Item = Option<bool>>, + ) -> Self { + let col = Column::from_name(name.into()); + + // take stats out and update them + let container_stats = self + .stats + .remove(&col) + .unwrap_or_default() + .with_contained(values, contained); + + // put stats back in + self.stats.insert(col, container_stats); + self + } + } + + impl PruningStatistics for TestStatistics { + fn min_values(&self, column: &Column) -> Option<ArrayRef> { + self.stats + .get(column) + .map(|container_stats| container_stats.min()) + .unwrap_or(None) + } + + fn max_values(&self, column: &Column) -> Option<ArrayRef> { + self.stats + .get(column) + .map(|container_stats| container_stats.max()) + .unwrap_or(None) + } + + fn num_containers(&self) -> usize { + self.stats + .values() + .next() + .map(|container_stats| container_stats.len()) + .unwrap_or(0) + } + + fn null_counts(&self, column: &Column) -> Option<ArrayRef> { + self.stats + .get(column) + .map(|container_stats| container_stats.null_counts()) + .unwrap_or(None) + } + + fn row_counts(&self) -> Option<ArrayRef> { + self.stats + .values() + .find_map(|container_stats| container_stats.row_counts()) + } + + fn contained( + &self, + column: &Column, + values: &HashSet<ScalarValue>, + ) -> Option<BooleanArray> { + self.stats + .get(column) + .and_then(|container_stats| container_stats.contained(values)) + } + } + + /// Returns the specified min/max container values + struct OneContainerStats { + min_values: Option<ArrayRef>, + max_values: Option<ArrayRef>, + num_containers: usize, + } + + impl PruningStatistics for OneContainerStats { + fn min_values(&self, _column: &Column) -> Option<ArrayRef> { + self.min_values.clone() + } + + fn max_values(&self, _column: &Column) -> Option<ArrayRef> { + self.max_values.clone() + } + + fn num_containers(&self) -> usize { + self.num_containers + } + + fn null_counts(&self, _column: &Column) -> Option<ArrayRef> { + None + } + + fn row_counts(&self) -> Option<ArrayRef> { + None + } + + fn contained( + &self, + _column: &Column, + _values: &HashSet<ScalarValue>, + ) -> Option<BooleanArray> { + None + } + } + + /// Row count should only be referenced once in the pruning expression, even if we need the row count + /// for multiple columns. + #[test] + fn test_unique_row_count_field_and_column() { + // c1 = 100 AND c2 = 200 + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, true), + Field::new("c2", DataType::Int32, true), + ])); + let expr = col("c1").eq(lit(100)).and(col("c2").eq(lit(200))); + let expr = logical2physical(&expr, &schema); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(expr) + .unwrap(); + // note pruning expression refers to row_count twice + assert_eq!( + "c1_null_count@2 != row_count@3 AND c1_min@0 <= 100 AND 100 <= c1_max@1 AND c2_null_count@6 != row_count@3 AND c2_min@4 <= 200 AND 200 <= c2_max@5", + p.predicate_expr.to_string() + ); + + // Fields in required schema should be unique, otherwise when creating batches + // it will fail because of duplicate field names + let mut fields = HashSet::new(); + for (_col, _ty, field) in p.required_columns().iter() { + let was_new = fields.insert(field); + if !was_new { + panic!( + "Duplicate field in required schema: {field:?}. Previous fields:\n{fields:#?}" + ); + } + } + } + + #[test] + fn prune_all_rows_null_counts() { + // if null_count = row_count then we should prune the container for i = 0 + // regardless of the statistics + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let statistics = TestStatistics::new().with( + "i", + ContainerStats::new_i32( + vec![Some(0)], // min + vec![Some(0)], // max + ) + .with_null_counts(vec![Some(1)]) + .with_row_counts(vec![Some(1)]), + ); + let expected_ret = &[false]; + prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret); + + // this should be true even if the container stats are missing + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let container_stats = ContainerStats { + min: Some(Arc::new(Int32Array::from(vec![None]))), + max: Some(Arc::new(Int32Array::from(vec![None]))), + null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))), + row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))), + ..ContainerStats::default() + }; + let statistics = TestStatistics::new().with("i", container_stats); + let expected_ret = &[false]; + prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret); + + // If the null counts themselves are missing we should be able to fall back to the stats + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let container_stats = ContainerStats { + min: Some(Arc::new(Int32Array::from(vec![Some(0)]))), + max: Some(Arc::new(Int32Array::from(vec![Some(0)]))), + null_counts: Some(Arc::new(UInt64Array::from(vec![None]))), + row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))), + ..ContainerStats::default() + }; + let statistics = TestStatistics::new().with("i", container_stats); + let expected_ret = &[true]; + prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret); + let expected_ret = &[false]; + prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret); + + // Same for the row counts + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let container_stats = ContainerStats { + min: Some(Arc::new(Int32Array::from(vec![Some(0)]))), + max: Some(Arc::new(Int32Array::from(vec![Some(0)]))), + null_counts: Some(Arc::new(UInt64Array::from(vec![Some(1)]))), + row_counts: Some(Arc::new(UInt64Array::from(vec![None]))), + ..ContainerStats::default() + }; + let statistics = TestStatistics::new().with("i", container_stats); + let expected_ret = &[true]; + prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret); + let expected_ret = &[false]; + prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret); + } + + #[test] + fn prune_missing_statistics() { + // If the min or max stats are missing we should not prune + // (unless we know all rows are null, see `prune_all_rows_null_counts`) + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let container_stats = ContainerStats { + min: Some(Arc::new(Int32Array::from(vec![None, Some(0)]))), + max: Some(Arc::new(Int32Array::from(vec![Some(0), None]))), + null_counts: Some(Arc::new(UInt64Array::from(vec![Some(0), Some(0)]))), + row_counts: Some(Arc::new(UInt64Array::from(vec![Some(1), Some(1)]))), + ..ContainerStats::default() + }; + let statistics = TestStatistics::new().with("i", container_stats); + let expected_ret = &[true, true]; + prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret); + let expected_ret = &[false, true]; + prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret); + let expected_ret = &[true, false]; + prune_with_expr(col("i").lt(lit(0)), &schema, &statistics, expected_ret); + } + + #[test] + fn prune_null_stats() { + // if null_count = row_count then we should prune the container for i = 0 + // regardless of the statistics + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + + let statistics = TestStatistics::new().with( + "i", + ContainerStats::new_i32( + vec![Some(0)], // min + vec![Some(0)], // max + ) + .with_null_counts(vec![Some(1)]) + .with_row_counts(vec![Some(1)]), + ); + + let expected_ret = &[false]; + + // i = 0 + prune_with_expr(col("i").eq(lit(0)), &schema, &statistics, expected_ret); + } + + #[test] + fn test_build_statistics_record_batch() { + // Request a record batch with of s1_min, s2_max, s3_max, s3_min + let required_columns = RequiredColumns::from(vec![ + // min of original column s1, named s1_min + ( + phys_expr::Column::new("s1", 1), + StatisticsType::Min, + Field::new("s1_min", DataType::Int32, true), + ), + // max of original column s2, named s2_max + ( + phys_expr::Column::new("s2", 2), + StatisticsType::Max, + Field::new("s2_max", DataType::Int32, true), + ), + // max of original column s3, named s3_max + ( + phys_expr::Column::new("s3", 3), + StatisticsType::Max, + Field::new("s3_max", DataType::Utf8, true), + ), + // min of original column s3, named s3_min + ( + phys_expr::Column::new("s3", 3), + StatisticsType::Min, + Field::new("s3_min", DataType::Utf8, true), + ), + ]); + + let statistics = TestStatistics::new() + .with( + "s1", + ContainerStats::new_i32( + vec![None, None, Some(9), None], // min + vec![Some(10), None, None, None], // max + ), + ) + .with( + "s2", + ContainerStats::new_i32( + vec![Some(2), None, None, None], // min + vec![Some(20), None, None, None], // max + ), + ) + .with( + "s3", + ContainerStats::new_utf8( + vec![Some("a"), None, None, None], // min + vec![Some("q"), None, Some("r"), None], // max + ), + ); + + let batch = + build_statistics_record_batch(&statistics, &required_columns).unwrap(); + assert_snapshot!(batches_to_string(&[batch]), @r" + +--------+--------+--------+--------+ + | s1_min | s2_max | s3_max | s3_min | + +--------+--------+--------+--------+ + | | 20 | q | a | + | | | | | + | 9 | | r | | + | | | | | + +--------+--------+--------+--------+ + "); + } + + #[test] + fn test_build_statistics_casting() { + // Test requesting a Timestamp column, but getting statistics as Int64 + // which is what Parquet does + + // Request a record batch with of s1_min as a timestamp + let required_columns = RequiredColumns::from(vec![( + phys_expr::Column::new("s3", 3), + StatisticsType::Min, + Field::new( + "s1_min", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ), + )]); + + // Note the statistics pass back i64 (not timestamp) + let statistics = OneContainerStats { + min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))), + max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))), + num_containers: 1, + }; + + let batch = + build_statistics_record_batch(&statistics, &required_columns).unwrap(); + + assert_snapshot!(batches_to_string(&[batch]), @r" + +-------------------------------+ + | s1_min | + +-------------------------------+ + | 1970-01-01T00:00:00.000000010 | + +-------------------------------+ + "); + } + + #[test] + fn test_build_statistics_no_required_stats() { + let required_columns = RequiredColumns::new(); + + let statistics = OneContainerStats { + min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))), + max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))), + num_containers: 1, + }; + + let batch = + build_statistics_record_batch(&statistics, &required_columns).unwrap(); + assert_eq!(batch.num_rows(), 1); // had 1 container + } + + #[test] + fn test_build_statistics_inconsistent_types() { + // Test requesting a Utf8 column when the stats return some other type + + // Request a record batch with of s1_min as a timestamp + let required_columns = RequiredColumns::from(vec![( + phys_expr::Column::new("s3", 3), + StatisticsType::Min, + Field::new("s1_min", DataType::Utf8, true), + )]); + + // Note the statistics return an invalid UTF-8 sequence which will be converted to null + let statistics = OneContainerStats { + min_values: Some(Arc::new(BinaryArray::from(vec![&[255u8] as &[u8]]))), + max_values: None, + num_containers: 1, + }; + + let batch = + build_statistics_record_batch(&statistics, &required_columns).unwrap(); + assert_snapshot!(batches_to_string(&[batch]), @r" + +--------+ + | s1_min | + +--------+ + | | + +--------+ + "); + } + + #[test] + fn test_build_statistics_inconsistent_length() { + // return an inconsistent length to the actual statistics arrays + let required_columns = RequiredColumns::from(vec![( + phys_expr::Column::new("s1", 3), + StatisticsType::Min, + Field::new("s1_min", DataType::Int64, true), + )]); + + // Note the statistics pass back i64 (not timestamp) + let statistics = OneContainerStats { + min_values: Some(Arc::new(Int64Array::from(vec![Some(10)]))), + max_values: Some(Arc::new(Int64Array::from(vec![Some(20)]))), + num_containers: 3, + }; + + let result = + build_statistics_record_batch(&statistics, &required_columns).unwrap_err(); + assert!( + result + .to_string() + .contains("mismatched statistics length. Expected 3, got 1"), + "{}", + result + ); + } + + #[test] + fn row_group_predicate_eq() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = + "c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1"; + + // test column on the left + let expr = col("c1").eq(lit(1)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(1).eq(col("c1")); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_not_eq() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = + "c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1)"; + + // test column on the left + let expr = col("c1").not_eq(lit(1)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(1).not_eq(col("c1")); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_gt() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 > 1"; + + // test column on the left + let expr = col("c1").gt(lit(1)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(1).lt(col("c1")); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_gt_eq() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_max@0 >= 1"; + + // test column on the left + let expr = col("c1").gt_eq(lit(1)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + // test column on the right + let expr = lit(1).lt_eq(col("c1")); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_lt() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1"; + + // test column on the left + let expr = col("c1").lt(lit(1)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(1).gt(col("c1")); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_lt_eq() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 <= 1"; + + // test column on the left + let expr = col("c1").lt_eq(lit(1)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + // test column on the right + let expr = lit(1).gt_eq(col("c1")); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_and() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + Field::new("c3", DataType::Int32, false), + ]); + // test AND operator joining supported c1 < 1 expression and unsupported c2 > c3 expression + let expr = col("c1").lt(lit(1)).and(col("c2").lt(col("c3"))); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_or() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // test OR operator joining supported c1 < 1 expression and unsupported c2 % 2 = 0 expression + let expr = col("c1").lt(lit(1)).or(col("c2").rem(lit(2)).eq(lit(0))); + let expected_expr = "true"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_not() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "true"; + + let expr = col("c1").not(); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_not_bool() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]); + let expected_expr = "NOT c1_min@0 AND c1_max@1"; + + let expr = col("c1").not(); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_bool() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]); + let expected_expr = "c1_min@0 OR c1_max@1"; + + let expr = col("c1"); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + /// Test that non-boolean literal expressions don't prune any containers and error gracefully by not pruning anything instead of e.g. panicking + #[test] + fn row_group_predicate_non_boolean() { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); + let statistics = TestStatistics::new() + .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)])); + let expected_ret = &[true]; + prune_with_expr(lit(1), &schema, &statistics, expected_ret); + } + + // Test that literal-to-literal comparisons are correctly evaluated. + // When both sides are constants, the expression should be evaluated directly + // and if it's false, all containers should be pruned. + #[test] + fn row_group_predicate_literal_false() { + // lit(1) = lit(2) is always false, so all containers should be pruned + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); + let statistics = TestStatistics::new() + .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)])); + let expected_ret = &[false]; + prune_with_simplified_expr(lit(1).eq(lit(2)), &schema, &statistics, expected_ret); + } + + /// Test nested/complex literal expression trees. + /// This is an integration test that PhysicalExprSimplifier + PruningPredicate work together as expected. + #[test] + fn row_group_predicate_literal_true() { + // lit(1) = lit(1) is always true, so no containers should be pruned + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); + let statistics = TestStatistics::new() + .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)])); + let expected_ret = &[true]; + prune_with_simplified_expr(lit(1).eq(lit(1)), &schema, &statistics, expected_ret); + } + + /// Test nested/complex literal expression trees. + /// This is an integration test that PhysicalExprSimplifier + PruningPredicate work together as expected. + #[test] + fn row_group_predicate_literal_null() { + // lit(1) = null is always null, so no containers should be pruned + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); + let statistics = TestStatistics::new() + .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)])); + let expected_ret = &[true]; + prune_with_simplified_expr( + lit(1).eq(lit(ScalarValue::Null)), + &schema, + &statistics, + expected_ret, + ); + } + + /// Test nested/complex literal expression trees. + /// This is an integration test that PhysicalExprSimplifier + PruningPredicate work together as expected. + #[test] + fn row_group_predicate_complex_literals() { + let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); + let statistics = TestStatistics::new() + .with("c1", ContainerStats::new_i32(vec![Some(0)], vec![Some(10)])); + + // (1 + 2) > 0 is always true + prune_with_simplified_expr( + (lit(1) + lit(2)).gt(lit(0)), + &schema, + &statistics, + &[true], + ); + + // (1 + 2) < 0 is always false + prune_with_simplified_expr( + (lit(1) + lit(2)).lt(lit(0)), + &schema, + &statistics, + &[false], + ); + + // Nested AND of literals: true AND false = false + prune_with_simplified_expr( + lit(true).and(lit(false)), + &schema, + &statistics, + &[false], + ); + + // Nested OR of literals: true OR false = true + prune_with_simplified_expr( + lit(true).or(lit(false)), + &schema, + &statistics, + &[true], + ); + + // Complex nested: (1 < 2) AND (3 > 1) = true AND true = true + prune_with_simplified_expr( + lit(1).lt(lit(2)).and(lit(3).gt(lit(1))), + &schema, + &statistics, + &[true], + ); + + // Complex nested: (1 > 2) OR (3 < 1) = false OR false = false + prune_with_simplified_expr( + lit(1).gt(lit(2)).or(lit(3).lt(lit(1))), + &schema, + &statistics, + &[false], + ); + } + + /// Integration test demonstrating that a dynamic filter with replaced children as literals will be snapshotted, simplified and then pruned correctly. + #[test] + fn row_group_predicate_dynamic_filter_with_literals() { + let schema = Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, true), + Field::new("part", DataType::Utf8, true), + ])); + let statistics = TestStatistics::new() + // Note that we have no stats, pruning can only happen via partition value pruning from the dynamic filter + .with_row_counts("c1", vec![Some(10)]); + let dynamic_filter_expr = col("c1").gt(lit(5)).and(col("part").eq(lit("B"))); + let phys_expr = logical2physical(&dynamic_filter_expr, &schema); + let children = collect_columns(&phys_expr) + .iter() + .map(|c| Arc::new(c.clone()) as Arc<dyn PhysicalExpr>) + .collect_vec(); + let dynamic_phys_expr = + Arc::new(DynamicFilterPhysicalExpr::new(children, phys_expr)) + as Arc<dyn PhysicalExpr>; + // Simulate the partition value substitution that would happen in ParquetOpener + let remapped_expr = dynamic_phys_expr + .children() + .into_iter() + .map(|child_expr| { + let Some(col_expr) = child_expr.downcast_ref::<phys_expr::Column>() + else { + return Arc::clone(child_expr); + }; + if col_expr.name() == "part" { + // simulate dynamic filter replacement with literal "A" + Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some( + "A".to_string(), + )))) as Arc<dyn PhysicalExpr> + } else { + Arc::clone(child_expr) + } + }) + .collect_vec(); + let dynamic_filter_expr = + dynamic_phys_expr.with_new_children(remapped_expr).unwrap(); + // After substitution the expression is c1 > 5 AND part = "B" which should prune the file since the partition value is "A" + let expected = &[false]; + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(dynamic_filter_expr) + .unwrap(); + let result = p.prune(&statistics).unwrap(); + assert_eq!(result, expected); + } + + #[test] + fn row_group_predicate_lt_bool() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Boolean, false)]); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < true"; + + // DF doesn't support arithmetic on boolean columns so + // this predicate will error when evaluated + let expr = col("c1").lt(lit(true)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_required_columns() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + let mut required_columns = RequiredColumns::new(); + // c1 < 1 and (c2 = 2 or c2 = 3) + let expr = col("c1") + .lt(lit(1)) + .and(col("c2").eq(lit(2)).or(col("c2").eq(lit(3)))); + let expected_expr = "c1_null_count@1 != row_count@2 AND c1_min@0 < 1 AND (c2_null_count@5 != row_count@2 AND c2_min@3 <= 2 AND 2 <= c2_max@4 OR c2_null_count@5 != row_count@2 AND c2_min@3 <= 3 AND 3 <= c2_max@4)"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut required_columns); + assert_eq!(predicate_expr.to_string(), expected_expr); + println!("required_columns: {required_columns:#?}"); // for debugging assertions below + // c1 < 1 should add c1_min + let c1_min_field = Field::new("c1_min", DataType::Int32, false); + assert_eq!( + required_columns.columns[0], + ( + phys_expr::Column::new("c1", 0), + StatisticsType::Min, + c1_min_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + // c1 < 1 should add c1_null_count + let c1_null_count_field = Field::new("c1_null_count", DataType::UInt64, false); + assert_eq!( + required_columns.columns[1], + ( + phys_expr::Column::new("c1", 0), + StatisticsType::NullCount, + c1_null_count_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + // c1 < 1 should add row_count + let row_count_field = Field::new("row_count", DataType::UInt64, false); + assert_eq!( + required_columns.columns[2], + ( + phys_expr::Column::new("c1", 0), + StatisticsType::RowCount, + row_count_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + // c2 = 2 should add c2_min and c2_max + let c2_min_field = Field::new("c2_min", DataType::Int32, false); + assert_eq!( + required_columns.columns[3], + ( + phys_expr::Column::new("c2", 1), + StatisticsType::Min, + c2_min_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + let c2_max_field = Field::new("c2_max", DataType::Int32, false); + assert_eq!( + required_columns.columns[4], + ( + phys_expr::Column::new("c2", 1), + StatisticsType::Max, + c2_max_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + // c2 = 2 should add c2_null_count + let c2_null_count_field = Field::new("c2_null_count", DataType::UInt64, false); + assert_eq!( + required_columns.columns[5], + ( + phys_expr::Column::new("c2", 1), + StatisticsType::NullCount, + c2_null_count_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + // c2 = 1 should add row_count + let row_count_field = Field::new("row_count", DataType::UInt64, false); + assert_eq!( + required_columns.columns[2], + ( + phys_expr::Column::new("c1", 0), + StatisticsType::RowCount, + row_count_field.with_nullable(true) // could be nullable if stats are not present + ) + ); + // c2 = 3 shouldn't add any new statistics fields + assert_eq!(required_columns.columns.len(), 6); + + Ok(()) + } + + #[test] + fn row_group_predicate_in_list() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // test c1 in(1, 2, 3) + let expr = Expr::InList(InList::new( + Box::new(col("c1")), + vec![lit(1), lit(2), lit(3)], + false, + )); + let expected_expr = "c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 2 AND 2 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 3 AND 3 <= c1_max@1"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_in_list_empty() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // test c1 in() + let expr = Expr::InList(InList::new(Box::new(col("c1")), vec![], false)); + let expected_expr = "true"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_in_list_negated() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // test c1 not in(1, 2, 3) + let expr = Expr::InList(InList::new( + Box::new(col("c1")), + vec![lit(1), lit(2), lit(3)], + true, + )); + let expected_expr = "c1_null_count@2 != row_count@3 AND (c1_min@0 != 1 OR 1 != c1_max@1) AND c1_null_count@2 != row_count@3 AND (c1_min@0 != 2 OR 2 != c1_max@1) AND c1_null_count@2 != row_count@3 AND (c1_min@0 != 3 OR 3 != c1_max@1)"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_between() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + + // test c1 BETWEEN 1 AND 5 + let expr1 = col("c1").between(lit(1), lit(5)); + + // test 1 <= c1 <= 5 + let expr2 = col("c1").gt_eq(lit(1)).and(col("c1").lt_eq(lit(5))); + + let predicate_expr1 = + test_build_predicate_expression(&expr1, &schema, &mut RequiredColumns::new()); + + let predicate_expr2 = + test_build_predicate_expression(&expr2, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr1.to_string(), predicate_expr2.to_string()); + + Ok(()) + } + + #[test] + fn row_group_predicate_between_with_in_list() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ]); + // test c1 in(1, 2) + let expr1 = col("c1").in_list(vec![lit(1), lit(2)], false); + + // test c2 BETWEEN 4 AND 5 + let expr2 = col("c2").between(lit(4), lit(5)); + + // test c1 in(1, 2) and c2 BETWEEN 4 AND 5 + let expr3 = expr1.and(expr2); + + let expected_expr = "(c1_null_count@2 != row_count@3 AND c1_min@0 <= 1 AND 1 <= c1_max@1 OR c1_null_count@2 != row_count@3 AND c1_min@0 <= 2 AND 2 <= c1_max@1) AND c2_null_count@5 != row_count@3 AND c2_max@4 >= 4 AND c2_null_count@5 != row_count@3 AND c2_min@6 <= 5"; + let predicate_expr = + test_build_predicate_expression(&expr3, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_in_list_to_many_values() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + // test c1 in(1..21) + // in pruning.rs has MAX_IN_LIST_SIZE = 20, more than this value will be rewrite + // always true + let expr = col("c1").in_list((1..=21).map(lit).collect(), false); + + let expected_expr = "true"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + // With the configurable cap, a caller that raises + // `max_in_list_size` above the default gets the IN list rewritten + // into a per-value min/max chain instead of falling through to `true`. + // This verifies both `PredicateRewriter::with_max_in_list_size` and the + // recursive OR path inside `build_predicate_expression`. + #[test] + fn row_group_predicate_in_list_rewritten_at_raised_cap() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + // 25 items — above the default 20, below a raised cap of 32. + let expr = col("c1").in_list((1..=25).map(lit).collect(), false); + let physical = logical2physical(&expr, &schema); + let rewriter = PredicateRewriter::new().with_max_in_list_size(32); + let predicate_expr = + rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema); + // At the raised cap, IN is rewritten into per-value min/max checks + // OR'd together; the resulting predicate must not collapse to + // `true` (which is what the default cap produces). + assert_ne!( + predicate_expr.to_string(), + "true", + "IN(25) with raised cap must rewrite into a statistics-based predicate, not fall through to `true`" + ); + // Sanity: the rewritten predicate references per-value literals. + assert!( + predicate_expr.to_string().contains(" <= 1 ") + && predicate_expr.to_string().contains(" <= 25 "), + "rewritten predicate should include per-value bounds for each IN entry, got: {predicate_expr}" + ); + Ok(()) + } + + // Guard: when the cap is 0 (opt-out) the IN branch is skipped entirely + // regardless of list length, so even a small IN falls through to the + // unhandled hook. + #[test] + fn row_group_predicate_in_list_disabled_at_zero_cap() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let expr = col("c1").in_list(vec![lit(1), lit(2), lit(3)], false); + let physical = logical2physical(&expr, &schema); + let rewriter = PredicateRewriter::new().with_max_in_list_size(0); + let predicate_expr = + rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema); + assert_eq!( + predicate_expr.to_string(), + "true", + "cap=0 must skip IN rewrite even for small lists" + ); + Ok(()) + } + + // The high-level [`PruningPredicateBuilder`] should thread + // `max_in_list_size` all the way through: a 25-item IN with the default + // cap must fall through to the unhandled hook (`predicate_expr = true`), + // while a raised cap produces a real per-value statistics predicate. + #[test] + fn pruning_predicate_builder_threads_max_in_list_size() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let expr = col("c1").in_list((1..=25).map(lit).collect(), false); + let physical = logical2physical(&expr, &schema); + + // With the default cap the IN branch bails out and the pruning + // predicate expression collapses to `true` (i.e., no container + // pruning based on stats). + let default_pp = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical))?; + assert_eq!( + default_pp.predicate_expr().to_string(), + "true", + "default cap must fall through to `true` for 25-item IN" + ); + + // Raising the cap produces a real statistics predicate with per- + // value bounds. + let raised_pp = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(32) + .try_build(physical)?; + let raised_expr = raised_pp.predicate_expr().to_string(); + assert_ne!( + raised_expr, "true", + "raised cap must produce a real statistics predicate for 25-item IN" + ); + assert!( + raised_expr.contains(" <= 1 ") && raised_expr.contains(" <= 25 "), + "raised-cap predicate should include per-value bounds, got: {raised_expr}" + ); + Ok(()) + } + + #[test] + #[expect(deprecated)] + fn deprecated_try_new_delegates_to_builder() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let expr = logical2physical(&col("c1").eq(lit(1)), &schema); + + let deprecated = + PruningPredicate::try_new(Arc::clone(&expr), Arc::clone(&schema))?; + let builder = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr)?; + + assert_eq!( + deprecated.predicate_expr().to_string(), + builder.predicate_expr().to_string() + ); + assert_eq!( + deprecated.required_columns().schema(), + builder.required_columns().schema() + ); + Ok(()) + } + + #[test] + fn row_group_predicate_cast_int_int() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64)"; + + // test cast(c1 as int64) = 1 + // test column on the left + let expr = cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(1)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Int64(Some(1))).eq(cast(col("c1"), DataType::Int64)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + let expected_expr = + "c1_null_count@1 != row_count@2 AND TRY_CAST(c1_max@0 AS Int64) > 1"; + + // test column on the left + let expr = + try_cast(col("c1"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(1)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = + lit(ScalarValue::Int64(Some(1))).lt(try_cast(col("c1"), DataType::Int64)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_string_string() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Utf8) <= 1 AND 1 <= CAST(c1_max@1 AS Utf8)"; + + // test column on the left + let expr = cast(col("c1"), DataType::Utf8) + .eq(lit(ScalarValue::Utf8(Some("1".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Utf8(Some("1".to_string()))) + .eq(cast(col("c1"), DataType::Utf8)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_string_int() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = cast(col("c1"), DataType::Int32).eq(lit(ScalarValue::Int32(Some(1)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Int32(Some(1))).eq(cast(col("c1"), DataType::Int32)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_int_string() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = cast(col("c1"), DataType::Utf8) + .eq(lit(ScalarValue::Utf8(Some("1".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Utf8(Some("1".to_string()))) + .eq(cast(col("c1"), DataType::Utf8)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_date_date() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Date64) <= 1970-01-01 AND 1970-01-01 <= CAST(c1_max@1 AS Date64)"; + + // test column on the left + let expr = + cast(col("c1"), DataType::Date64).eq(lit(ScalarValue::Date64(Some(123)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = + lit(ScalarValue::Date64(Some(123))).eq(cast(col("c1"), DataType::Date64)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_dict_string_date() -> Result<()> { + // Test with Dictionary<UInt8, Utf8> for the literal + let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = cast( + col("c1"), + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + ) + .eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))).eq(cast( + col("c1"), + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + )); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_date_dict_string() -> Result<()> { + // Test with Dictionary<UInt8, Utf8> for the column + let schema = Schema::new(vec![Field::new( + "c1", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + false, + )]); + let expected_expr = "true"; + + // test column on the left + let expr = + cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = + lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_dict_dict_same_value_type() -> Result<()> { + // Test with Dictionary types that have the same value type but different key types + let schema = Schema::new(vec![Field::new( + "c1", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)), + false, + )]); + + // Direct comparison with no cast + let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + let expected_expr = + "c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1"; + assert_eq!(predicate_expr.to_string(), expected_expr); + + // Test with column cast to a dictionary with different key type + let expr = cast( + col("c1"), + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)), + ) + .eq(lit(ScalarValue::Utf8(Some("test".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Dictionary(UInt16, Utf8)) <= test AND test <= CAST(c1_max@1 AS Dictionary(UInt16, Utf8))"; + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_dict_dict_different_value_type() -> Result<()> { + // Test with Dictionary types that have different value types + let schema = Schema::new(vec![Field::new( + "c1", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Int32)), + false, + )]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 123 AND 123 <= CAST(c1_max@1 AS Int64)"; + + // Test with literal of a different type + let expr = + cast(col("c1"), DataType::Int64).eq(lit(ScalarValue::Int64(Some(123)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_nested_dict() -> Result<()> { + // Test with nested Dictionary types + let schema = Schema::new(vec![Field::new( + "c1", + DataType::Dictionary( + Box::new(DataType::UInt8), + Box::new(DataType::Dictionary( + Box::new(DataType::UInt16), + Box::new(DataType::Utf8), + )), + ), + false, + )]); + let expected_expr = + "c1_null_count@2 != row_count@3 AND c1_min@0 <= test AND test <= c1_max@1"; + + // Test with a simple literal + let expr = col("c1").eq(lit(ScalarValue::Utf8(Some("test".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_dict_date_dict_date() -> Result<()> { + // Test with dictionary-wrapped date types for both sides + let schema = Schema::new(vec![Field::new( + "c1", + DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Date32)), + false, + )]); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Dictionary(UInt16, Date64)) <= 1970-01-01 AND 1970-01-01 <= CAST(c1_max@1 AS Dictionary(UInt16, Date64))"; + + // Test with a cast to a different date type + let expr = cast( + col("c1"), + DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Date64)), + ) + .eq(lit(ScalarValue::Date64(Some(123)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_date_string() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = + cast(col("c1"), DataType::Date32).eq(lit(ScalarValue::Date32(Some(123)))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = + lit(ScalarValue::Date32(Some(123))).eq(cast(col("c1"), DataType::Date32)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_string_date() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Date32, false)]); + let expected_expr = "true"; + + // test column on the left + let expr = cast(col("c1"), DataType::Utf8) + .eq(lit(ScalarValue::Utf8(Some("2024-01-01".to_string())))); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + // test column on the right + let expr = lit(ScalarValue::Utf8(Some("2024-01-01".to_string()))) + .eq(cast(col("c1"), DataType::Utf8)); + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn row_group_predicate_cast_list() -> Result<()> { + let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); + // test cast(c1 as int64) in int64(1, 2, 3) + let expr = Expr::InList(InList::new( + Box::new(cast(col("c1"), DataType::Int64)), + vec![ + lit(ScalarValue::Int64(Some(1))), + lit(ScalarValue::Int64(Some(2))), + lit(ScalarValue::Int64(Some(3))), + ], + false, + )); + let expected_expr = "c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 1 AND 1 <= CAST(c1_max@1 AS Int64) OR c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 2 AND 2 <= CAST(c1_max@1 AS Int64) OR c1_null_count@2 != row_count@3 AND CAST(c1_min@0 AS Int64) <= 3 AND 3 <= CAST(c1_max@1 AS Int64)"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + let expr = Expr::InList(InList::new( + Box::new(cast(col("c1"), DataType::Int64)), + vec![ + lit(ScalarValue::Int64(Some(1))), + lit(ScalarValue::Int64(Some(2))), + lit(ScalarValue::Int64(Some(3))), + ], + true, + )); + let expected_expr = "c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 1 OR 1 != CAST(c1_max@1 AS Int64)) AND c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 2 OR 2 != CAST(c1_max@1 AS Int64)) AND c1_null_count@2 != row_count@3 AND (CAST(c1_min@0 AS Int64) != 3 OR 3 != CAST(c1_max@1 AS Int64))"; + let predicate_expr = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + assert_eq!(predicate_expr.to_string(), expected_expr); + + Ok(()) + } + + #[test] + fn prune_decimal_data() { + // decimal(9,2) + let schema = Arc::new(Schema::new(vec![Field::new( + "s1", + DataType::Decimal128(9, 2), + true, + )])); + + prune_with_expr( + // s1 > 5 + col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 9, 2))), + &schema, + // If the data is written by spark, the physical data type is INT32 in the parquet + // So we use the INT32 type of statistic. + &TestStatistics::new().with( + "s1", + ContainerStats::new_i32( + vec![Some(0), Some(4), None, Some(3)], // min + vec![Some(5), Some(6), Some(4), None], // max + ), + ), + &[false, true, false, true], + ); + + prune_with_expr( + // with cast column to other type + cast(col("s1"), DataType::Decimal128(14, 3)) + .gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))), + &schema, + &TestStatistics::new().with( + "s1", + ContainerStats::new_i32( + vec![Some(0), Some(4), None, Some(3)], // min + vec![Some(5), Some(6), Some(4), None], // max + ), + ), + &[false, true, false, true], + ); + + prune_with_expr( + // with try cast column to other type + try_cast(col("s1"), DataType::Decimal128(14, 3)) + .gt(lit(ScalarValue::Decimal128(Some(5000), 14, 3))), + &schema, + &TestStatistics::new().with( + "s1", + ContainerStats::new_i32( + vec![Some(0), Some(4), None, Some(3)], // min + vec![Some(5), Some(6), Some(4), None], // max + ), + ), + &[false, true, false, true], + ); + + // decimal(18,2) + let schema = Arc::new(Schema::new(vec![Field::new( + "s1", + DataType::Decimal128(18, 2), + true, + )])); + prune_with_expr( + // s1 > 5 + col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 18, 2))), + &schema, + // If the data is written by spark, the physical data type is INT64 in the parquet + // So we use the INT32 type of statistic. + &TestStatistics::new().with( + "s1", + ContainerStats::new_i64( + vec![Some(0), Some(4), None, Some(3)], // min + vec![Some(5), Some(6), Some(4), None], // max + ), + ), + &[false, true, false, true], + ); + + // decimal(23,2) + let schema = Arc::new(Schema::new(vec![Field::new( + "s1", + DataType::Decimal128(23, 2), + true, + )])); + + prune_with_expr( + // s1 > 5 + col("s1").gt(lit(ScalarValue::Decimal128(Some(500), 23, 2))), + &schema, + &TestStatistics::new().with( + "s1", + ContainerStats::new_decimal128( + vec![Some(0), Some(400), None, Some(300)], // min + vec![Some(500), Some(600), Some(400), None], // max + 23, + 2, + ), + ), + &[false, true, false, true], + ); + } + + #[test] + fn prune_api() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s1", DataType::Utf8, true), + Field::new("s2", DataType::Int32, true), + ])); + + let statistics = TestStatistics::new().with( + "s2", + ContainerStats::new_i32( + vec![Some(0), Some(4), None, Some(3)], // min + vec![Some(5), Some(6), None, None], // max + ), + ); + prune_with_expr( + // Prune using s2 > 5 + col("s2").gt(lit(5)), + &schema, + &statistics, + // s2 [0, 5] ==> no rows should pass + // s2 [4, 6] ==> some rows could pass + // No stats for s2 ==> some rows could pass + // s2 [3, None] (null max) ==> some rows could pass + &[false, true, true, true], + ); + + prune_with_expr( + // filter with cast + cast(col("s2"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(5)))), + &schema, + &statistics, + &[false, true, true, true], + ); + } + + #[test] + fn prune_not_eq_data() { + let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)])); + + prune_with_expr( + // Prune using s2 != 'M' + col("s1").not_eq(lit("M")), + &schema, + &TestStatistics::new().with( + "s1", + ContainerStats::new_utf8( + vec![Some("A"), Some("A"), Some("N"), Some("M"), None, Some("A")], // min + vec![Some("Z"), Some("L"), Some("Z"), Some("M"), None, None], // max + ), + ), + // s1 [A, Z] ==> might have values that pass predicate + // s1 [A, L] ==> all rows pass the predicate + // s1 [N, Z] ==> all rows pass the predicate + // s1 [M, M] ==> all rows do not pass the predicate + // No stats for s2 ==> some rows could pass + // s2 [3, None] (null max) ==> some rows could pass + &[true, true, true, false, true, true], + ); + } + + /// Creates setup for boolean chunk pruning + /// + /// For predicate "b1" (boolean expr) + /// b1 [false, false] ==> no rows can pass (not keep) + /// b1 [false, true] ==> some rows could pass (must keep) + /// b1 [true, true] ==> all rows must pass (must keep) + /// b1 [NULL, NULL] ==> unknown (must keep) + /// b1 [false, NULL] ==> unknown (must keep) + /// + /// For predicate "!b1" (boolean expr) + /// b1 [false, false] ==> all rows pass (must keep) + /// b1 [false, true] ==> some rows could pass (must keep) + /// b1 [true, true] ==> no rows can pass (not keep) + /// b1 [NULL, NULL] ==> unknown (must keep) + /// b1 [false, NULL] ==> unknown (must keep) + fn bool_setup() -> (SchemaRef, TestStatistics, Vec<bool>, Vec<bool>) { + let schema = + Arc::new(Schema::new(vec![Field::new("b1", DataType::Boolean, true)])); + + let statistics = TestStatistics::new().with( + "b1", + ContainerStats::new_bool( + vec![Some(false), Some(false), Some(true), None, Some(false)], // min + vec![Some(false), Some(true), Some(true), None, None], // max + ), + ); + let expected_true = vec![false, true, true, true, true]; + let expected_false = vec![true, true, false, true, true]; + + (schema, statistics, expected_true, expected_false) + } + + #[test] + fn prune_bool_const_expr() { + let (schema, statistics, _, _) = bool_setup(); + + prune_with_expr( + // true + lit(true), + &schema, + &statistics, + &[true, true, true, true, true], + ); + + prune_with_expr( + // false + lit(false), + &schema, + &statistics, + &[false, false, false, false, false], + ); + } + + #[test] + fn prune_bool_column() { + let (schema, statistics, expected_true, _) = bool_setup(); + + prune_with_expr( + // b1 + col("b1"), + &schema, + &statistics, + &expected_true, + ); + } + + #[test] + fn prune_bool_not_column() { + let (schema, statistics, _, expected_false) = bool_setup(); + + prune_with_expr( + // !b1 + col("b1").not(), + &schema, + &statistics, + &expected_false, + ); + } + + #[test] + fn prune_bool_column_eq_true() { + let (schema, statistics, expected_true, _) = bool_setup(); + + prune_with_expr( + // b1 = true + col("b1").eq(lit(true)), + &schema, + &statistics, + &expected_true, + ); + } + + #[test] + fn prune_bool_not_column_eq_true() { + let (schema, statistics, _, expected_false) = bool_setup(); + + prune_with_expr( + // !b1 = true + col("b1").not().eq(lit(true)), + &schema, + &statistics, + &expected_false, + ); + } + + /// Creates a setup for chunk pruning, modeling a int32 column "i" + /// with 5 different containers (e.g. RowGroups). They have [min, + /// max]: + /// + /// i [-5, 5] + /// i [1, 11] + /// i [-11, -1] + /// i [NULL, NULL] + /// i [1, NULL] + fn int32_setup() -> (SchemaRef, TestStatistics) { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + + let statistics = TestStatistics::new().with( + "i", + ContainerStats::new_i32( + vec![Some(-5), Some(1), Some(-11), None, Some(1)], // min + vec![Some(5), Some(11), Some(-1), None, None], // max + ), + ); + (schema, statistics) + } + + #[test] + fn prune_int32_col_gt_zero() { + let (schema, statistics) = int32_setup(); + + // Expression "i > 0" and "-i < 0" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> all rows must pass (must keep) + // i [-11, -1] ==> no rows can pass (not keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> unknown (must keep) + let expected_ret = &[true, true, false, true, true]; + + // i > 0 + prune_with_expr(col("i").gt(lit(0)), &schema, &statistics, expected_ret); + + // -i < 0 + prune_with_expr( + Expr::Negative(Box::new(col("i"))).lt(lit(0)), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_lte_zero() { + let (schema, statistics) = int32_setup(); + + // Expression "i <= 0" and "-i >= 0" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> no rows can pass (not keep) + // i [-11, -1] ==> all rows must pass (must keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> no rows can pass (not keep) + let expected_ret = &[true, false, true, true, false]; + + prune_with_expr( + // i <= 0 + col("i").lt_eq(lit(0)), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // -i >= 0 + Expr::Negative(Box::new(col("i"))).gt_eq(lit(0)), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_lte_zero_cast() { + let (schema, statistics) = int32_setup(); + + // Expression "cast(i as utf8) <= '0'" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> no rows can pass in theory, -0.22 (conservatively keep) + // i [-11, -1] ==> no rows could pass in theory (conservatively keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> no rows can pass (conservatively keep) + let expected_ret = &[true, true, true, true, true]; + + prune_with_expr( + // cast(i as utf8) <= 0 + cast(col("i"), DataType::Utf8).lt_eq(lit("0")), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // try_cast(i as utf8) <= 0 + try_cast(col("i"), DataType::Utf8).lt_eq(lit("0")), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // cast(-i as utf8) >= 0 + cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // try_cast(-i as utf8) >= 0 + try_cast(Expr::Negative(Box::new(col("i"))), DataType::Utf8).gt_eq(lit("0")), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_eq_zero() { + let (schema, statistics) = int32_setup(); + + // Expression "i = 0" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> no rows can pass (not keep) + // i [-11, -1] ==> no rows can pass (not keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> no rows can pass (not keep) + let expected_ret = &[true, false, false, true, false]; + + prune_with_expr( + // i = 0 + col("i").eq(lit(0)), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_is_not_distinct_from() { + let (schema, statistics) = int32_setup(); + + // Without null counts, IS NOT DISTINCT FROM a non-null literal can + // still use min/max ranges, but unknown all-null containers must be kept. + let expected_ret = &[true, false, false, true, false]; + + prune_with_expr( + is_not_distinct_from(col("i"), lit(0)), + &schema, + &statistics, + expected_ret, + ); + + // The operator is symmetric, so the scalar-left form should prune the + // same row groups. + prune_with_expr( + is_not_distinct_from(lit(0), col("i")), + &schema, + &statistics, + expected_ret, + ); + + let statistics = statistics + .with_row_counts("i", vec![Some(10), Some(9), None, Some(4), Some(10)]) + .with_null_counts("i", vec![Some(0), Some(1), None, Some(4), Some(0)]); + + let expected_ret = &[true, false, false, false, false]; + prune_with_expr( + is_not_distinct_from(col("i"), lit(0)), + &schema, + &statistics, + expected_ret, + ); + + let expected_ret = &[false, true, true, true, false]; + prune_with_expr( + is_not_distinct_from(col("i"), lit(ScalarValue::Int32(None))), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_is_distinct_from() { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let statistics = TestStatistics::new().with( + "i", + ContainerStats::new_i32( + vec![Some(0), Some(0), Some(5), None], + vec![Some(0), Some(2), Some(5), None], + ) + .with_row_counts(vec![Some(2), Some(2), Some(2), Some(2)]) + .with_null_counts(vec![Some(0), Some(0), Some(0), Some(2)]), + ); + + let expected_ret = &[false, true, true, true]; + prune_with_expr( + is_distinct_from(col("i"), lit(0)), + &schema, + &statistics, + expected_ret, + ); + + // The operator is symmetric, so the scalar-left form should prune the + // same row groups. + prune_with_expr( + is_distinct_from(lit(0), col("i")), + &schema, + &statistics, + expected_ret, + ); + + let expected_ret = &[true, true, true, false]; + prune_with_expr( + is_distinct_from(col("i"), lit(ScalarValue::Int32(None))), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_eq_zero_cast() { + let (schema, statistics) = int32_setup(); + + // Expression "cast(i as int64) = 0" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> no rows can pass (not keep) + // i [-11, -1] ==> no rows can pass (not keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> no rows can pass (not keep) + let expected_ret = &[true, false, false, true, false]; + + prune_with_expr( + cast(col("i"), DataType::Int64).eq(lit(0i64)), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + try_cast(col("i"), DataType::Int64).eq(lit(0i64)), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_eq_zero_cast_as_str() { + let (schema, statistics) = int32_setup(); + + // Note the cast is to a string where sorting properties are + // not the same as integers + // + // Expression "cast(i as utf8) = '0'" + // i [-5, 5] ==> some rows could pass (keep) + // i [1, 11] ==> no rows can pass (could keep) + // i [-11, -1] ==> no rows can pass (could keep) + // i [NULL, NULL] ==> unknown (keep) + // i [1, NULL] ==> no rows can pass (could keep) + let expected_ret = &[true, true, true, true, true]; + + prune_with_expr( + cast(col("i"), DataType::Utf8).eq(lit("0")), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_col_lt_neg_one() { + let (schema, statistics) = int32_setup(); + + // Expression "i > -1" and "-i < 1" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> all rows must pass (must keep) + // i [-11, -1] ==> no rows can pass (not keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> all rows must pass (must keep) + let expected_ret = &[true, true, false, true, true]; + + prune_with_expr( + // i > -1 + col("i").gt(lit(-1)), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // -i < 1 + Expr::Negative(Box::new(col("i"))).lt(lit(1)), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_is_null() { + let (schema, statistics) = int32_setup(); + + // Expression "i IS NULL" when there are no null statistics, + // should all be kept + let expected_ret = &[true, true, true, true, true]; + + prune_with_expr( + // i IS NULL, no null statistics + col("i").is_null(), + &schema, + &statistics, + expected_ret, + ); + + // provide null counts for each column + let statistics = statistics.with_null_counts( + "i", + vec![ + Some(0), // no nulls (don't keep) + Some(1), // 1 null + None, // unknown nulls + None, // unknown nulls (min/max are both null too, like no stats at all) + Some(0), // 0 nulls (max=null too which means no known max) (don't keep) + ], + ); + + let expected_ret = &[false, true, true, true, false]; + + prune_with_expr( + // i IS NULL, with actual null statistics + col("i").is_null(), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_int32_column_is_known_all_null() { + let (schema, statistics) = int32_setup(); + + // Expression "i < 0" + // i [-5, 5] ==> some rows could pass (must keep) + // i [1, 11] ==> no rows can pass (not keep) + // i [-11, -1] ==> all rows must pass (must keep) + // i [NULL, NULL] ==> unknown (must keep) + // i [1, NULL] ==> no rows can pass (not keep) + let expected_ret = &[true, false, true, true, false]; + + prune_with_expr( + // i < 0 + col("i").lt(lit(0)), + &schema, + &statistics, + expected_ret, + ); + + // provide row counts for each column + let statistics = statistics.with_row_counts( + "i", + vec![ + Some(10), // 10 rows of data + Some(9), // 9 rows of data + None, // unknown row counts + Some(4), + Some(10), + ], + ); + + // pruning result is still the same if we only know row counts + prune_with_expr( + // i < 0, with only row counts statistics + col("i").lt(lit(0)), + &schema, + &statistics, + expected_ret, + ); + + // provide null counts for each column + let statistics = statistics.with_null_counts( + "i", + vec![ + Some(0), // no nulls + Some(1), // 1 null + None, // unknown nulls + Some(4), // 4 nulls, which is the same as the row counts, i.e. this column is all null (don't keep) + Some(0), // 0 nulls (max=null too which means no known max) + ], + ); + + // Expression "i < 0" with actual null and row counts statistics + // col | min, max | row counts | null counts | + // ----+--------------+------------+-------------+ + // i | [-5, 5] | 10 | 0 | ==> Some rows could pass (must keep) + // i | [1, 11] | 9 | 1 | ==> No rows can pass (not keep) + // i | [-11,-1] | Unknown | Unknown | ==> All rows must pass (must keep) + // i | [NULL, NULL] | 4 | 4 | ==> The column is all null (not keep) + // i | [1, NULL] | 10 | 0 | ==> No rows can pass (not keep) + let expected_ret = &[true, false, true, false, false]; + + prune_with_expr( + // i < 0, with actual null and row counts statistics + col("i").lt(lit(0)), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn prune_cast_scalar() { + // The data type of column i is INT32 + let (schema, statistics) = int32_setup(); + let expected_ret = &[true, true, false, true, true]; + + prune_with_expr( + // i > int64(0) + col("i").gt(cast(lit(ScalarValue::Int64(Some(0))), DataType::Int32)), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // cast(i as int64) > int64(0) + cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // try_cast(i as int64) > int64(0) + try_cast(col("i"), DataType::Int64).gt(lit(ScalarValue::Int64(Some(0)))), + &schema, + &statistics, + expected_ret, + ); + + prune_with_expr( + // `-cast(i as int64) < 0` convert to `cast(i as int64) > -0` + Expr::Negative(Box::new(cast(col("i"), DataType::Int64))) + .lt(lit(ScalarValue::Int64(Some(0)))), + &schema, + &statistics, + expected_ret, + ); + } + + #[test] + fn test_increment_utf8() { + // Basic ASCII + assert_eq!(increment_utf8("abc").unwrap(), "abd"); + assert_eq!(increment_utf8("abz").unwrap(), "ab{"); + + // Test around ASCII 127 (DEL) + assert_eq!(increment_utf8("~").unwrap(), "\u{7f}"); // 126 -> 127 + assert_eq!(increment_utf8("\u{7f}").unwrap(), "\u{80}"); // 127 -> 128 + + // Test 2-byte UTF-8 sequences + assert_eq!(increment_utf8("ß").unwrap(), "à"); // U+00DF -> U+00E0 + + // Test 3-byte UTF-8 sequences + assert_eq!(increment_utf8("℣").unwrap(), "ℤ"); // U+2123 -> U+2124 + + // Test at UTF-8 boundaries + assert_eq!(increment_utf8("\u{7FF}").unwrap(), "\u{800}"); // 2-byte to 3-byte boundary + assert_eq!(increment_utf8("\u{FFFF}").unwrap(), "\u{10000}"); // 3-byte to 4-byte boundary + + // Test that if we can't increment we return None + assert!(increment_utf8("").is_none()); + assert!(increment_utf8("\u{10FFFF}").is_none()); // U+10FFFF is the max code point + + // Test that if we can't increment the last character we do the previous one and truncate + assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b"); + + // Test surrogate pair range (0xD800..=0xDFFF) + assert_eq!(increment_utf8("a\u{D7FF}").unwrap(), "b"); + assert!(increment_utf8("\u{D7FF}").is_none()); + + // Test non-characters range (0xFDD0..=0xFDEF) + assert_eq!(increment_utf8("a\u{FDCF}").unwrap(), "b"); + assert!(increment_utf8("\u{FDCF}").is_none()); + + // Test private use area limit (>= 0x110000) + assert_eq!(increment_utf8("a\u{10FFFF}").unwrap(), "b"); + assert!(increment_utf8("\u{10FFFF}").is_none()); // Can't increment past max valid codepoint + } + + /// Creates a setup for chunk pruning, modeling a utf8 column "s1" + /// with 5 different containers (e.g. RowGroups). They have [min, + /// max]: + /// s1 ["A", "Z"] + /// s1 ["A", "L"] + /// s1 ["N", "Z"] + /// s1 [NULL, NULL] + /// s1 ["A", NULL] + /// s1 ["", "A"] + /// s1 ["", ""] + /// s1 ["AB", "A\u{10ffff}"] + /// s1 ["A\u{10ffff}\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] + fn utf8_setup() -> (SchemaRef, TestStatistics) { + let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)])); + + let statistics = TestStatistics::new().with( + "s1", + ContainerStats::new_utf8( + vec![ + Some("A"), + Some("A"), + Some("N"), + Some("M"), + None, + Some("A"), + Some(""), + Some(""), + Some("AB"), + Some("A\u{10ffff}\u{10ffff}"), + ], // min + vec![ + Some("Z"), + Some("L"), + Some("Z"), + Some("M"), + None, + None, + Some("A"), + Some(""), + Some("A\u{10ffff}\u{10ffff}\u{10ffff}"), + Some("A\u{10ffff}\u{10ffff}"), + ], // max + ), + ); + (schema, statistics) + } + + #[test] + fn prune_utf8_eq() { + let (schema, statistics) = utf8_setup(); + + let expr = col("s1").eq(lit("A")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["M", "M"] ==> no rows can pass (not keep) + false, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> no rows can pass (not keep) + false, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").eq(lit("")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["A", "L"] ==> no rows can pass (not keep) + false, + // s1 ["N", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["M", "M"] ==> no rows can pass (not keep) + false, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> no rows can pass (not keep) + false, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> all rows must pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + #[test] + fn prune_utf8_not_eq() { + let (schema, statistics) = utf8_setup(); + + let expr = col("s1").not_eq(lit("A")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["M", "M"] ==> all rows must pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> all rows must pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").not_eq(lit("")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["A", "L"] ==> all rows must pass (must keep) + true, + // s1 ["N", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["M", "M"] ==> all rows must pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> no rows can pass (not keep) + false, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + #[test] + fn prune_utf8_like_one() { + let (schema, statistics) = utf8_setup(); + + let expr = col("s1").like(lit("A_")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["M", "M"] ==> no rows can pass (not keep) + false, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> no rows can pass (not keep) + false, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit("_A_")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["M", "M"] ==> some rows could pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> some rows could pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit("_")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["A", "L"] ==> all rows must pass (must keep) + true, + // s1 ["N", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["M", "M"] ==> all rows must pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> all rows must pass (must keep) + true, + // s1 ["", ""] ==> all rows must pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit("")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["A", "L"] ==> no rows can pass (not keep) + false, + // s1 ["N", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["M", "M"] ==> no rows can pass (not keep) + false, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> no rows can pass (not keep) + false, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> all rows must pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + #[test] + fn prune_utf8_like_many() { + let (schema, statistics) = utf8_setup(); + + let expr = col("s1").like(lit("A%")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["M", "M"] ==> no rows can pass (not keep) + false, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> no rows can pass (not keep) + false, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit("%A%")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["M", "M"] ==> some rows could pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> some rows could pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit("%")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["A", "L"] ==> all rows must pass (must keep) + true, + // s1 ["N", "Z"] ==> all rows must pass (must keep) + true, + // s1 ["M", "M"] ==> all rows must pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> unknown (must keep) + true, + // s1 ["", "A"] ==> all rows must pass (must keep) + true, + // s1 ["", ""] ==> all rows must pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> all rows must pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit("")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["A", "L"] ==> no rows can pass (not keep) + false, + // s1 ["N", "Z"] ==> no rows can pass (not keep) + false, + // s1 ["M", "M"] ==> no rows can pass (not keep) + false, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> no rows can pass (not keep) + false, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> all rows must pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + // `build_like_match()` must honor `\` escapes when scanning the pattern for + // wildcards. + #[test] + fn prune_utf8_like_escaped_chars() { + let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)])); + let statistics = TestStatistics::new().with( + "s1", + ContainerStats::new_utf8( + vec![ + Some("foo_aaa"), + Some(r#"foo\aaa"#), + Some("foo"), + Some("bar"), + Some("foo%aaa"), + Some("%foo_aaa"), + ], // min + vec![ + Some("foo_zzz"), + Some(r#"foo\zzz"#), + Some("foozzz"), + Some("baz"), + Some("foo%zzz"), + Some("%foo_zzz"), + ], // max + ), + ); + + let expr = col("s1").like(lit(r#"foo\_%"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => every value starts with literal + // "foo_" and matches the pattern; must keep. + true, + // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep) + false, + // s1 ["foo", "foozzz"] => stats don't prove "foo_" is or isn't in + // range; must conservatively keep. + true, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep) + false, + // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit(r#"foo\\%"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep) + false, + // s1 ["foo\aaa", "foo\zzz"] => every value starts with literal + // "foo\" and matches the pattern; must keep. + true, + // s1 ["foo", "foozzz"] => stats don't prove "foo\" is or isn't in + // range; must conservatively keep. + true, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep) + false, + // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").like(lit(r#"foo\%%"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep) + false, + // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep) + false, + // s1 ["foo", "foozzz"] => range straddles "foo%"; must keep. + true, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => every value starts with literal + // "foo%" and matches the pattern; must keep. + true, + // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + // No wildcard after escapes: pattern reduces to an equality check on + // the literal "foo_". + let expr = col("s1").like(lit(r#"foo\_"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep) + false, + // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep) + false, + // s1 ["foo", "foozzz"] => "foo_" is within the range; must keep. + true, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep) + false, + // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + // Leading escaped `%`: prefix is "%foo" (non-empty), so the guard + // for "all wildcards" must NOT bail out here. + let expr = col("s1").like(lit(r#"\%foo%"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep) + false, + // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep) + false, + // s1 ["foo", "foozzz"] => no rows can pass (not keep) + false, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep) + false, + // s1 ["%foo_aaa", "%foo_zzz"] => every value starts with literal + // "%foo" and matches the pattern; must keep. + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + // Two escaped wildcards, no real wildcard: equality on "foo%_". + let expr = col("s1").like(lit(r#"foo\%\_"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep) + false, + // s1 ["foo\aaa", "foo\zzz"] => no rows can pass (not keep) + false, + // s1 ["foo", "foozzz"] => "foo%_" is within the range; must keep. + true, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep) + false, + // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + // Escaped backslash followed by more literal chars before the + // wildcard: prefix is "foo\bar". + let expr = col("s1").like(lit(r#"foo\\bar%"#)); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["foo_aaa", "foo_zzz"] => no rows can pass (not keep) + false, + // s1 ["foo\aaa", "foo\zzz"] => range straddles "foo\bar"; must + // keep. + true, + // s1 ["foo", "foozzz"] => range straddles "foo\bar"; must keep. + true, + // s1 ["bar", "baz"] => no rows can pass (not keep) + false, + // s1 ["foo%aaa", "foo%zzz"] => no rows can pass (not keep) + false, + // s1 ["%foo_aaa", "%foo_zzz"] => no rows can pass (not keep) + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + #[test] + fn prune_utf8_not_like_one() { + let (schema, statistics) = utf8_setup(); + + let expr = col("s1").not_like(lit("A\u{10ffff}_")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["M", "M"] ==> some rows could pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> some rows could pass (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> some rows could pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no row match. (min, max) maybe truncate + // original (min, max) maybe ("A\u{10ffff}\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}\u{10ffff}\u{10ffff}") + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + #[test] + fn prune_utf8_not_like_many() { + let (schema, statistics) = utf8_setup(); + + let expr = col("s1").not_like(lit("A\u{10ffff}%")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["M", "M"] ==> some rows could pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> some rows could pass (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> some rows could pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> no row match + false, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["M", "M"] ==> some rows could pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> some rows could pass (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> some rows could pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").not_like(lit("A\u{10ffff}%\u{10ffff}_")); + #[rustfmt::skip] + let expected_ret = &[ + // s1 ["A", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["A", "L"] ==> some rows could pass (must keep) + true, + // s1 ["N", "Z"] ==> some rows could pass (must keep) + true, + // s1 ["M", "M"] ==> some rows could pass (must keep) + true, + // s1 [NULL, NULL] ==> unknown (must keep) + true, + // s1 ["A", NULL] ==> some rows could pass (must keep) + true, + // s1 ["", "A"] ==> some rows could pass (must keep) + true, + // s1 ["", ""] ==> some rows could pass (must keep) + true, + // s1 ["AB", "A\u{10ffff}\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + // s1 ["A\u{10ffff}\u{10ffff}", "A\u{10ffff}\u{10ffff}"] ==> some rows could pass (must keep) + true, + ]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + + let expr = col("s1").not_like(lit("A\\%%")); + let statistics = TestStatistics::new().with( + "s1", + ContainerStats::new_utf8( + vec![Some("A%a"), Some("A")], + vec![Some("A%c"), Some("A")], + ), + ); + let expected_ret = &[false, true]; + prune_with_expr(expr, &schema, &statistics, expected_ret); + } + + #[test] + fn test_rewrite_expr_to_prunable() { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let df_schema = DFSchema::try_from(schema.clone()).unwrap(); + + // column op lit + let left_input = col("a"); + let left_input = logical2physical(&left_input, &schema); + let right_input = lit(ScalarValue::Int32(Some(12))); + let right_input = logical2physical(&right_input, &schema); + let (result_left, _, result_right) = rewrite_expr_to_prunable( + &left_input, + Operator::Eq, + &right_input, + df_schema.clone(), + ) + .unwrap(); + assert_eq!(result_left.to_string(), left_input.to_string()); + assert_eq!(result_right.to_string(), right_input.to_string()); + + // cast op lit + let left_input = cast(col("a"), DataType::Decimal128(20, 3)); + let left_input = logical2physical(&left_input, &schema); + let right_input = lit(ScalarValue::Decimal128(Some(12), 20, 3)); + let right_input = logical2physical(&right_input, &schema); + let (result_left, _, result_right) = rewrite_expr_to_prunable( + &left_input, + Operator::Gt, + &right_input, + df_schema.clone(), + ) + .unwrap(); + assert_eq!(result_left.to_string(), left_input.to_string()); + assert_eq!(result_right.to_string(), right_input.to_string()); + + // try_cast op lit + let left_input = try_cast(col("a"), DataType::Int64); + let left_input = logical2physical(&left_input, &schema); + let right_input = lit(ScalarValue::Int64(Some(12))); + let right_input = logical2physical(&right_input, &schema); + let (result_left, _, result_right) = + rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema) + .unwrap(); + assert_eq!(result_left.to_string(), left_input.to_string()); + assert_eq!(result_right.to_string(), right_input.to_string()); + + // TODO: add test for other case and op + } + + #[test] + fn test_rewrite_expr_to_prunable_custom_unhandled_hook() { + struct CustomUnhandledHook; + + impl UnhandledPredicateHook for CustomUnhandledHook { + /// This handles an arbitrary case of a column that doesn't exist in the schema + /// by renaming it to yet another column that doesn't exist in the schema + /// (the transformation is arbitrary, the point is that it can do whatever it wants) + fn handle(&self, _expr: &Arc<dyn PhysicalExpr>) -> Arc<dyn PhysicalExpr> { + Arc::new(phys_expr::Literal::new(ScalarValue::Int32(Some(42)))) + } + } + + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let schema_with_b = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]); + + let rewriter = PredicateRewriter::new() + .with_unhandled_hook(Arc::new(CustomUnhandledHook {})); + + let transform_expr = |expr| { + let expr = logical2physical(&expr, &schema_with_b); + rewriter.rewrite_predicate_to_statistics_predicate(&expr, &schema) + }; + + // transform an arbitrary valid expression that we know is handled + let known_expression = col("a").eq(lit(12)); + let known_expression_transformed = PredicateRewriter::new() + .rewrite_predicate_to_statistics_predicate( + &logical2physical(&known_expression, &schema), + &schema, + ); + + // an expression referencing an unknown column (that is not in the schema) gets passed to the hook + let input = col("b").eq(lit(12)); + let expected = logical2physical(&lit(42), &schema); + let transformed = transform_expr(input.clone()); + assert_eq!(transformed.to_string(), expected.to_string()); + + // more complex case with unknown column + let input = known_expression.clone().and(input.clone()); + let expected = phys_expr::BinaryExpr::new( + Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed), + Operator::And, + logical2physical(&lit(42), &schema), + ); + let transformed = transform_expr(input.clone()); + assert_eq!(transformed.to_string(), expected.to_string()); + + // an unknown expression gets passed to the hook + let input = array_has(make_array(vec![lit(1)]), col("a")); + let expected = logical2physical(&lit(42), &schema); + let transformed = transform_expr(input.clone()); + assert_eq!(transformed.to_string(), expected.to_string()); + + // more complex case with unknown expression + let input = known_expression.and(input); + let expected = phys_expr::BinaryExpr::new( + Arc::<dyn PhysicalExpr>::clone(&known_expression_transformed), + Operator::And, + logical2physical(&lit(42), &schema), + ); + let transformed = transform_expr(input.clone()); + assert_eq!(transformed.to_string(), expected.to_string()); + } + + #[test] + fn test_rewrite_expr_to_prunable_error() { + // cast string value to numeric value + // this cast is not supported + let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]); + let df_schema = DFSchema::try_from(schema.clone()).unwrap(); + let left_input = cast(col("a"), DataType::Int64); + let left_input = logical2physical(&left_input, &schema); + let right_input = lit(ScalarValue::Int64(Some(12))); + let right_input = logical2physical(&right_input, &schema); + let result = rewrite_expr_to_prunable( + &left_input, + Operator::Gt, + &right_input, + df_schema.clone(), + ); + assert!(result.is_err()); + + // other expr + let left_input = is_null(col("a")); + let left_input = logical2physical(&left_input, &schema); + let right_input = lit(ScalarValue::Int64(Some(12))); + let right_input = logical2physical(&right_input, &schema); + let result = + rewrite_expr_to_prunable(&left_input, Operator::Gt, &right_input, df_schema); + assert!(result.is_err()); + // TODO: add other negative test for other case and op + } + + #[test] + fn prune_with_contained_one_column() { + let schema = Arc::new(Schema::new(vec![Field::new("s1", DataType::Utf8, true)])); + + // Model having information like a bloom filter for s1 + let statistics = TestStatistics::new() + .with_contained( + "s1", + [ScalarValue::from("foo")], + [ + // container 0 known to only contain "foo"", + Some(true), + // container 1 known to not contain "foo" + Some(false), + // container 2 unknown about "foo" + None, + // container 3 known to only contain "foo" + Some(true), + // container 4 known to not contain "foo" + Some(false), + // container 5 unknown about "foo" + None, + // container 6 known to only contain "foo" + Some(true), + // container 7 known to not contain "foo" + Some(false), + // container 8 unknown about "foo" + None, + ], + ) + .with_contained( + "s1", + [ScalarValue::from("bar")], + [ + // containers 0,1,2 known to only contain "bar" + Some(true), + Some(true), + Some(true), + // container 3,4,5 known to not contain "bar" + Some(false), + Some(false), + Some(false), + // container 6,7,8 unknown about "bar" + None, + None, + None, + ], + ) + .with_contained( + // the way the tests are setup, this data is + // consulted if the "foo" and "bar" are being checked at the same time + "s1", + [ScalarValue::from("foo"), ScalarValue::from("bar")], + [ + // container 0,1,2 unknown about ("foo, "bar") + None, + None, + None, + // container 3,4,5 known to contain only either "foo" and "bar" + Some(true), + Some(true), + Some(true), + // container 6,7,8 known to contain neither "foo" and "bar" + Some(false), + Some(false), + Some(false), + ], + ); + + // s1 = 'foo' + prune_with_expr( + col("s1").eq(lit("foo")), + &schema, + &statistics, + // rule out containers ('false) where we know foo is not present + &[true, false, true, true, false, true, true, false, true], + ); + + // s1 = 'bar' + prune_with_expr( + col("s1").eq(lit("bar")), + &schema, + &statistics, + // rule out containers where we know bar is not present + &[true, true, true, false, false, false, true, true, true], + ); + + // s1 = 'baz' (unknown value) + prune_with_expr( + col("s1").eq(lit("baz")), + &schema, + &statistics, + // can't rule out anything + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 = 'foo' AND s1 = 'bar' + prune_with_expr( + col("s1").eq(lit("foo")).and(col("s1").eq(lit("bar"))), + &schema, + &statistics, + // logically this predicate can't possibly be true (the column can't + // take on both values) but we could rule it out if the stats tell + // us that both values are not present + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 = 'foo' OR s1 = 'bar' + prune_with_expr( + col("s1").eq(lit("foo")).or(col("s1").eq(lit("bar"))), + &schema, + &statistics, + // can rule out containers that we know contain neither foo nor bar + &[true, true, true, true, true, true, false, false, false], + ); + + // s1 = 'foo' OR s1 = 'baz' + prune_with_expr( + col("s1").eq(lit("foo")).or(col("s1").eq(lit("baz"))), + &schema, + &statistics, + // can't rule out anything container + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 = 'foo' OR s1 = 'bar' OR s1 = 'baz' + prune_with_expr( + col("s1") + .eq(lit("foo")) + .or(col("s1").eq(lit("bar"))) + .or(col("s1").eq(lit("baz"))), + &schema, + &statistics, + // can rule out any containers based on knowledge of s1 and `foo`, + // `bar` and (`foo`, `bar`) + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 != foo + prune_with_expr( + col("s1").not_eq(lit("foo")), + &schema, + &statistics, + // rule out containers we know for sure only contain foo + &[false, true, true, false, true, true, false, true, true], + ); + + // s1 != bar + prune_with_expr( + col("s1").not_eq(lit("bar")), + &schema, + &statistics, + // rule out when we know for sure s1 has the value bar + &[false, false, false, true, true, true, true, true, true], + ); + + // s1 != foo AND s1 != bar + prune_with_expr( + col("s1") + .not_eq(lit("foo")) + .and(col("s1").not_eq(lit("bar"))), + &schema, + &statistics, + // can rule out any container where we know s1 does not have either 'foo' or 'bar' + &[true, true, true, false, false, false, true, true, true], + ); + + // s1 != foo AND s1 != bar AND s1 != baz + prune_with_expr( + col("s1") + .not_eq(lit("foo")) + .and(col("s1").not_eq(lit("bar"))) + .and(col("s1").not_eq(lit("baz"))), + &schema, + &statistics, + // can't rule out any container based on knowledge of s1,s2 + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 != foo OR s1 != bar + prune_with_expr( + col("s1") + .not_eq(lit("foo")) + .or(col("s1").not_eq(lit("bar"))), + &schema, + &statistics, + // cant' rule out anything based on contains information + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 != foo OR s1 != bar OR s1 != baz + prune_with_expr( + col("s1") + .not_eq(lit("foo")) + .or(col("s1").not_eq(lit("bar"))) + .or(col("s1").not_eq(lit("baz"))), + &schema, + &statistics, + // cant' rule out anything based on contains information + &[true, true, true, true, true, true, true, true, true], + ); + } + + #[test] + fn prune_with_contained_two_columns() { + let schema = Arc::new(Schema::new(vec![ + Field::new("s1", DataType::Utf8, true), + Field::new("s2", DataType::Utf8, true), + ])); + + // Model having information like bloom filters for s1 and s2 + let statistics = TestStatistics::new() + .with_contained( + "s1", + [ScalarValue::from("foo")], + [ + // container 0, s1 known to only contain "foo"", + Some(true), + // container 1, s1 known to not contain "foo" + Some(false), + // container 2, s1 unknown about "foo" + None, + // container 3, s1 known to only contain "foo" + Some(true), + // container 4, s1 known to not contain "foo" + Some(false), + // container 5, s1 unknown about "foo" + None, + // container 6, s1 known to only contain "foo" + Some(true), + // container 7, s1 known to not contain "foo" + Some(false), + // container 8, s1 unknown about "foo" + None, + ], + ) + .with_contained( + "s2", // for column s2 + [ScalarValue::from("bar")], + [ + // containers 0,1,2 s2 known to only contain "bar" + Some(true), + Some(true), + Some(true), + // container 3,4,5 s2 known to not contain "bar" + Some(false), + Some(false), + Some(false), + // container 6,7,8 s2 unknown about "bar" + None, + None, + None, + ], + ); + + // s1 = 'foo' + prune_with_expr( + col("s1").eq(lit("foo")), + &schema, + &statistics, + // rule out containers where we know s1 is not present + &[true, false, true, true, false, true, true, false, true], + ); + + // s1 = 'foo' OR s2 = 'bar' + let expr = col("s1").eq(lit("foo")).or(col("s2").eq(lit("bar"))); + prune_with_expr( + expr, + &schema, + &statistics, + // can't rule out any container (would need to prove that s1 != foo AND s2 != bar) + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 = 'foo' AND s2 != 'bar' + prune_with_expr( + col("s1").eq(lit("foo")).and(col("s2").not_eq(lit("bar"))), + &schema, + &statistics, + // can only rule out container where we know either: + // 1. s1 doesn't have the value 'foo` or + // 2. s2 has only the value of 'bar' + &[false, false, false, true, false, true, true, false, true], + ); + + // s1 != 'foo' AND s2 != 'bar' + prune_with_expr( + col("s1") + .not_eq(lit("foo")) + .and(col("s2").not_eq(lit("bar"))), + &schema, + &statistics, + // Can rule out any container where we know either + // 1. s1 has only the value 'foo' + // 2. s2 has only the value 'bar' + &[false, false, false, false, true, true, false, true, true], + ); + + // s1 != 'foo' AND (s2 = 'bar' OR s2 = 'baz') + prune_with_expr( + col("s1") + .not_eq(lit("foo")) + .and(col("s2").eq(lit("bar")).or(col("s2").eq(lit("baz")))), + &schema, + &statistics, + // Can rule out any container where we know s1 has only the value + // 'foo'. Can't use knowledge of s2 and bar to rule out anything + &[false, true, true, false, true, true, false, true, true], + ); + + // s1 like '%foo%bar%' + prune_with_expr( + col("s1").like(lit("foo%bar%")), + &schema, + &statistics, + // cant rule out anything with information we know + &[true, true, true, true, true, true, true, true, true], + ); + + // s1 like '%foo%bar%' AND s2 = 'bar' + prune_with_expr( + col("s1") + .like(lit("foo%bar%")) + .and(col("s2").eq(lit("bar"))), + &schema, + &statistics, + // can rule out any container where we know s2 does not have the value 'bar' + &[true, true, true, false, false, false, true, true, true], + ); + + // s1 like '%foo%bar%' OR s2 = 'bar' + prune_with_expr( + col("s1").like(lit("foo%bar%")).or(col("s2").eq(lit("bar"))), + &schema, + &statistics, + // can't rule out anything (we would have to prove that both the + // like and the equality must be false) + &[true, true, true, true, true, true, true, true, true], + ); + } + + #[test] + fn prune_with_range_and_contained() { + // Setup mimics range information for i, a bloom filter for s + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, true), + Field::new("s", DataType::Utf8, true), + ])); + + let statistics = TestStatistics::new() + .with( + "i", + ContainerStats::new_i32( + // Container 0, 3, 6: [-5 to 5] + // Container 1, 4, 7: [10 to 20] + // Container 2, 5, 9: unknown + vec![ + Some(-5), + Some(10), + None, + Some(-5), + Some(10), + None, + Some(-5), + Some(10), + None, + ], // min + vec![ + Some(5), + Some(20), + None, + Some(5), + Some(20), + None, + Some(5), + Some(20), + None, + ], // max + ), + ) + // Add contained information about the s and "foo" + .with_contained( + "s", + [ScalarValue::from("foo")], + [ + // container 0,1,2 known to only contain "foo" + Some(true), + Some(true), + Some(true), + // container 3,4,5 known to not contain "foo" + Some(false), + Some(false), + Some(false), + // container 6,7,8 unknown about "foo" + None, + None, + None, + ], + ); + + // i = 0 and s = 'foo' + prune_with_expr( + col("i").eq(lit(0)).and(col("s").eq(lit("foo"))), + &schema, + &statistics, + // Can rule out container where we know that either: + // 1. 0 is outside the min/max range of i + // 1. s does not contain foo + // (range is false, and contained is false) + &[true, false, true, false, false, false, true, false, true], + ); + + // i = 0 and s != 'foo' + prune_with_expr( + col("i").eq(lit(0)).and(col("s").not_eq(lit("foo"))), + &schema, + &statistics, + // Can rule out containers where either: + // 1. 0 is outside the min/max range of i + // 2. s only contains foo + &[false, false, false, true, false, true, true, false, true], + ); + + // i = 0 OR s = 'foo' + prune_with_expr( + col("i").eq(lit(0)).or(col("s").eq(lit("foo"))), + &schema, + &statistics, + // in theory could rule out containers if we had min/max values for + // s as well. But in this case we don't so we can't rule out anything + &[true, true, true, true, true, true, true, true, true], + ); + } + + /// prunes the specified expr with the specified schema and statistics, and + /// ensures it returns expected. + /// + /// `expected` is a vector of bools, where true means the row group should + /// be kept, and false means it should be pruned. + // TODO refactor other tests to use this to reduce boiler plate + fn prune_with_expr( + expr: Expr, + schema: &SchemaRef, + statistics: &TestStatistics, + expected: &[bool], + ) { + println!("Pruning with expr: {expr}"); + let expr = logical2physical(&expr, schema); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::<Schema>::clone(schema)) + .try_build(expr) + .unwrap(); + let result = p.prune(statistics).unwrap(); + assert_eq!(result, expected); + } + + fn prune_with_simplified_expr( + expr: Expr, + schema: &SchemaRef, + statistics: &TestStatistics, + expected: &[bool], + ) { + println!("Pruning with expr: {expr}"); + let expr = logical2physical(&expr, schema); + let simplifier = PhysicalExprSimplifier::new(schema); + let expr = simplifier.simplify(expr).unwrap(); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::<Schema>::clone(schema)) + .try_build(expr) + .unwrap(); + let result = p.prune(statistics).unwrap(); + assert_eq!(result, expected); + } + + fn is_not_distinct_from(left: Expr, right: Expr) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(left), + Operator::IsNotDistinctFrom, + Box::new(right), + )) + } + + fn is_distinct_from(left: Expr, right: Expr) -> Expr { + Expr::BinaryExpr(BinaryExpr::new( + Box::new(left), + Operator::IsDistinctFrom, + Box::new(right), + )) + } + + fn test_build_predicate_expression( + expr: &Expr, + schema: &Schema, + required_columns: &mut RequiredColumns, + ) -> Arc<dyn PhysicalExpr> { + let expr = logical2physical(expr, schema); + let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _; + build_predicate_expression( + &expr, + &Arc::new(schema.clone()), + required_columns, + &unhandled_hook, + MAX_IN_LIST_SIZE, + ) + } + + #[test] + fn test_build_predicate_expression_with_false() { + let expr = lit(ScalarValue::Boolean(Some(false))); + let schema = Schema::empty(); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + let expected = logical2physical(&expr, &schema); + assert_eq!(&res, &expected); + } + + #[test] + fn test_build_predicate_expression_with_and_false() { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); + let expr = and( + col("c1").eq(lit("a")), + lit(ScalarValue::Boolean(Some(false))), + ); + let res = + test_build_predicate_expression(&expr, &schema, &mut RequiredColumns::new()); + let expected = logical2physical(&lit(ScalarValue::Boolean(Some(false))), &schema); + assert_eq!(&res, &expected); + } + + #[test] + fn test_build_predicate_expression_with_or_false() { + let schema = Schema::new(vec![Field::new("c1", DataType::Utf8View, false)]); + let left_expr = col("c1").eq(lit("a")); + let right_expr = lit(ScalarValue::Boolean(Some(false))); + let res = test_build_predicate_expression( + &or(left_expr.clone(), right_expr.clone()), + &schema, + &mut RequiredColumns::new(), + ); + let expected = + "c1_null_count@2 != row_count@3 AND c1_min@0 <= a AND a <= c1_max@1"; + assert_eq!(res.to_string(), expected); + } + + /// A hash-table lookup over `column` carrying a bitmap over `keys`. + fn hash_lookup( + column: &Arc<dyn PhysicalExpr>, + keys: &[i64], + ) -> Arc<dyn PhysicalExpr> { + let array: ArrayRef = Arc::new(Int64Array::from(keys.to_vec())); + let bitmap = KeyRangeBitmap::try_new( + &array, + &ScalarValue::Int64(Some(*keys.iter().min().unwrap())), + &ScalarValue::Int64(Some(*keys.iter().max().unwrap())), + keys.len(), + ); + Arc::new( + HashTableLookupExpr::new( + vec![Arc::clone(column)], + SeededRandomState::with_seed(1), + Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(1)))), + "hash_lookup".to_string(), + ) + .with_pruning_bitmap(bitmap.map(Arc::new)), + ) + } + + #[test] + fn test_hash_lookup_pruning_via_min_max() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); + // Gapped keys: 10,20,30 then a far-away 100000, so the bitmap is sparse. + let lookup = hash_lookup(&column, &[10, 20, 30, 100_000]); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(lookup) + .unwrap(); + + // An opaque lookup yields no literal guarantees, so `contained()` is never + // consulted and every exclusion below comes from the min/max rewrite. + assert!(predicate.literal_guarantees().is_empty()); + + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i64( + vec![Some(5000), Some(15), Some(50000)], + vec![Some(8000), Some(25), Some(60000)], + ), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 + // holds 20. The bitmap excludes the first and third. + assert_eq!(result, vec![false, true, false]); + } + + #[test] + fn test_hash_lookup_pruning_with_nulls_and_missing_statistics() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); + let predicate = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(hash_lookup(&column, &[10, 20, 30, 100_000])) + .unwrap(); + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i64( + vec![Some(15), None, None, None, Some(5000), Some(5000)], + vec![Some(25), None, None, Some(25), None, Some(8000)], + ) + .with_row_counts(vec![Some(2); 6]) + .with_null_counts(vec![ + Some(0), + Some(2), + Some(0), + Some(0), + Some(0), + None, + ]), + ); + + // A matching range is kept and an all-null group is rejected. Missing + // either bound cannot prove absence. A known gap still excludes a group + // when the null count is unknown: neither nulls nor non-nulls can match. + assert_eq!( + predicate.prune(&statistics).unwrap(), + vec![true, false, true, true, true, false] + ); + } + + #[test] + fn test_case_without_else_cannot_prove_full_match() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, false)])); + let condition = logical2physical(&col("b").eq(lit(1)), &schema); + let case: Arc<dyn PhysicalExpr> = Arc::new( + phys_expr::CaseExpr::try_new( + None, + vec![(condition, phys_expr::lit(false))], + None, + ) + .unwrap(), + ); + // NOT(CASE WHEN b = 1 THEN false END) yields NULL for b=0, even + // though the referenced column is non-nullable. Such a row must not be + // counted as a match when a Parquet scan tries to satisfy a LIMIT. + let predicate: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::NotExpr::new(case)); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![0, 1]))], + ) + .unwrap(); + assert_eq!( + predicate + .evaluate(&batch) + .unwrap() + .into_array(2) + .unwrap() + .as_boolean(), + &BooleanArray::from(vec![None, Some(true)]) + ); + + // Full-match inference negates the predicate, simplifies double NOT, + // and asks if the result prunes the group. Dropping the implicit NULL + // branch would produce false and incorrectly mark every row matched. + let inverse = PhysicalExprSimplifier::new(&schema) + .simplify(Arc::new(phys_expr::NotExpr::new(predicate))) + .unwrap(); + let inverse = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(inverse) + .unwrap(); + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i32(vec![Some(0)], vec![Some(1)]) + .with_row_counts(vec![Some(2)]) + .with_null_counts(vec![Some(0)]), + ); + assert_eq!(inverse.prune(&statistics).unwrap(), vec![true]); + } + + #[rstest] + #[case::widening_i32_to_i64(DataType::Int32, DataType::Int64, true)] + #[case::widening_u32_to_i64(DataType::UInt32, DataType::Int64, true)] + #[case::widening_u16_to_i64(DataType::UInt16, DataType::Int64, true)] + #[case::same_type_i64_to_i64(DataType::Int64, DataType::Int64, true)] + #[case::narrowing_i64_to_i32(DataType::Int64, DataType::Int32, false)] + #[case::narrowing_u32_to_i16(DataType::UInt32, DataType::Int16, false)] + fn test_hash_lookup_pruning_through_cast( + #[case] file_type: DataType, + #[case] cast_type: DataType, + #[case] pruned: bool, + ) { + let schema = + Arc::new(Schema::new(vec![Field::new("b", file_type.clone(), true)])); + let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); + let cast_column: Arc<dyn PhysicalExpr> = + Arc::new(phys_expr::CastExpr::new(column, cast_type, None)); + let lookup = hash_lookup(&cast_column, &[10, 20, 30, 100_000]); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(lookup) + .unwrap(); + + let to_file_type = |values: Vec<Option<i64>>| -> ArrayRef { + let values: ArrayRef = Arc::new(values.into_iter().collect::<Int64Array>()); + arrow::compute::cast(&values, &file_type).unwrap() + }; + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new() + .with_min(to_file_type(vec![Some(5000), Some(15), Some(50000)])) + .with_max(to_file_type(vec![Some(8000), Some(25), Some(60000)])), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Containers 0 and 2 sit in the gap between 30 and 100000; container 1 + // holds 20. A working cast excludes the first and third; a declined one + // falls back to keeping every container. + let expected = if pruned { + vec![false, true, false] + } else { + vec![true, true, true] + }; + assert_eq!(result, expected); + } + + #[test] + fn test_partition_routed_hash_lookup_pruning() { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); + let column: Arc<dyn PhysicalExpr> = Arc::new(phys_expr::Column::new("b", 0)); + let literal = |value: ScalarValue| -> Arc<dyn PhysicalExpr> { + Arc::new(phys_expr::Literal::new(value)) + }; + // One branch per build partition, each holding its own slice of the build + // side, as `build_partitioned_filter` produces them. + let case: Arc<dyn PhysicalExpr> = Arc::new( + phys_expr::CaseExpr::try_new( + Some(literal(ScalarValue::UInt64(Some(0)))), + vec![ + ( + literal(ScalarValue::UInt64(Some(0))), + hash_lookup(&column, &[10, 20, 30]), + ), + ( + literal(ScalarValue::UInt64(Some(1))), + hash_lookup(&column, &[40, 50, 60]), + ), + ], + // Partitions with no build rows reject everything routed to them. + Some(literal(ScalarValue::Boolean(Some(false)))), + ) + .unwrap(), + ); + + let predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(case) + .unwrap(); + + let statistics = TestStatistics::new().with( + "b", + ContainerStats::new_i32( + vec![Some(5), Some(15), Some(45), Some(100)], + vec![Some(8), Some(25), Some(55), Some(200)], + ), + ); + + let result = predicate.prune(&statistics).unwrap(); + // Container 1 ([15,25]) holds 20 from the first branch and container 2 + // ([45,55]) holds 50 from the second, so a branch may match in each. The + // other two intersect neither branch, nor the `ELSE`. + assert_eq!(result, vec![false, true, true, false]); + } +}
diff --git a/native/vendor/rustfmt.toml b/native/vendor/rustfmt.toml new file mode 100644 index 0000000..c680d9d --- /dev/null +++ b/native/vendor/rustfmt.toml
@@ -0,0 +1,24 @@ +# 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. + +edition = "2024" +max_width = 90 + +# ignore generated files +# ignore = [ +# "arrow/src/ipc/gen", +#]