diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 60cff658a6a97..f1d7dc703ee96 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -252,6 +252,11 @@ harness = false name = "parquet_struct_projection" required-features = ["parquet"] +[[bench]] +harness = false +name = "cse_projection_pushdown" +required-features = ["parquet"] + [[bench]] harness = false name = "range_and_generate_series" diff --git a/datafusion/core/benches/cse_projection_pushdown.rs b/datafusion/core/benches/cse_projection_pushdown.rs new file mode 100644 index 0000000000000..f5f9ec55e8912 --- /dev/null +++ b/datafusion/core/benches/cse_projection_pushdown.rs @@ -0,0 +1,184 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the interaction between Common Subexpression Elimination +//! (CSE) and projection pushdown on parquet sources. +//! +//! Each query repeats a scalar function call several times, which the logical +//! CSE pass extracts into a single intermediate projection referenced by +//! column. These benchmarks measure the end-to-end cost of such queries, which +//! is dominated by how many times the extracted expression is ultimately +//! evaluated per row. + +use arrow::array::{Float64Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::instant::Instant; +use futures::stream::StreamExt; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use rand::prelude::*; +use rand::rng; +use std::sync::Arc; +use tempfile::NamedTempFile; + +const NUM_BATCHES: usize = 1024; +const BATCH_SIZE: usize = 1024; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + Field::new("c", DataType::Int64, false), + ])) +} + +fn generate_batch() -> RecordBatch { + let mut rng = rng(); + let len = BATCH_SIZE; + + let a: Float64Array = (0..len) + .map(|_| Some(rng.random_range(1.0..1000.0))) + .collect(); + let b: Float64Array = (0..len) + .map(|_| Some(rng.random_range(1.0..1000.0))) + .collect(); + let c: Int64Array = (0..len) + .map(|_| Some(rng.random_range(1i64..1000))) + .collect(); + + RecordBatch::try_new(schema(), vec![Arc::new(a), Arc::new(b), Arc::new(c)]).unwrap() +} + +fn generate_file() -> NamedTempFile { + let now = Instant::now(); + let mut named_file = tempfile::Builder::new() + .prefix("cse_projection_pushdown") + .suffix(".parquet") + .tempfile() + .unwrap(); + + println!("Generating parquet file - {}", named_file.path().display()); + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_max_row_group_row_count(Some(1024 * 1024)) + .build(); + + let mut writer = + ArrowWriter::try_new(&mut named_file, schema(), Some(props)).unwrap(); + + for _ in 0..NUM_BATCHES { + let batch = generate_batch(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + println!( + "Generated parquet file in {} seconds", + now.elapsed().as_secs_f32() + ); + + named_file +} + +fn criterion_benchmark(c: &mut Criterion) { + let temp_file = generate_file(); + let file_path = temp_file.path().display().to_string(); + + let partitions = 4; + let config = SessionConfig::new().with_target_partitions(partitions); + let context = SessionContext::new_with_config(config); + + let local_rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + + let query_rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(partitions) + .build() + .unwrap(); + + local_rt + .block_on(context.register_parquet("t", file_path.as_str(), Default::default())) + .unwrap(); + + // Queries that repeat a scalar function call, which CSE extracts into a + // single intermediate projection referenced by column. + let queries = vec![ + // Same sqrt(a) appears 3 times. + ( + "repeated_sqrt", + "SELECT sqrt(a) + 1, sqrt(a) * 2, sqrt(a) / b FROM t", + ), + // power(a, 2) appears in multiple places. + ( + "repeated_power", + "SELECT power(a, 2) + b, power(a, 2) - b, power(a, 2) * c FROM t", + ), + // Deeper nesting: ln(abs(a)) repeated. + ( + "repeated_nested_fn", + "SELECT ln(abs(a)) + 1, ln(abs(a)) * b, ln(abs(a)) + c FROM t", + ), + // Mixed: some repeated, some unique. + ( + "mixed_repeated_unique", + "SELECT sqrt(a) + sqrt(a), abs(b), sqrt(a) * c FROM t", + ), + // A trivial function (abs) repeated. + ( + "repeated_cheap_abs", + "SELECT abs(a) + 1, abs(a) * 2, abs(a) / b FROM t", + ), + // Baseline: no repeated expressions (CSE does not fire). + ( + "no_repeated_exprs", + "SELECT sqrt(a), abs(b), power(a, 2) FROM t", + ), + ]; + + for (name, query) in queries { + c.bench_function(&format!("cse_pushdown: {name}"), |b| { + b.iter(|| { + let query = query.to_string(); + let context = context.clone(); + let (sender, mut receiver) = futures::channel::mpsc::unbounded(); + + query_rt.spawn(async move { + let query = context.sql(&query).await.unwrap(); + let mut stream = query.execute_stream().await.unwrap(); + + while let Some(next) = stream.next().await { + sender.unbounded_send(next).unwrap(); + } + }); + + local_rt.block_on(async { + while receiver.next().await.transpose().unwrap().is_some() {} + }) + }); + }); + } + + drop(temp_file); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 1336ee69cd6dd..0d4586ce2fbaa 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -625,23 +625,29 @@ fn project_output_partitioning( } } -/// Returns `true` if merging `outer` into `inner` would duplicate a volatile -/// expression; the caller should then decline the merge. +/// Returns `true` if merging `outer` into `inner` would duplicate a volatile or +/// non-trivial expression that CSE deduplicated; the caller should then decline +/// the merge. /// -/// `inner` is the scan's current projection and `outer` the projection being -/// pushed into it; merging substitutes each `inner` expression into every -/// `outer` reference to it. If a volatile `inner` expression (e.g. `random()`, -/// `uuid()`) is referenced more than once, that single value gets inlined at -/// each site and re-evaluated independently, so references meant to share a -/// "locked-in" value diverge. This is the volatility guard the physical -/// `ProjectionPushdown` and `FilterPushdown` rules already apply (see -/// `datafusion_physical_expr_common::physical_expr::is_volatile`). +/// Merging substitutes each `inner` expression into every `outer` reference to +/// it. Since the logical optimizer extracts a repeated expression into a single +/// `inner` entry referenced by column, re-inlining it at more than one +/// reference site undoes that deduplication. An `inner` expression referenced +/// more than once is therefore blocked when it is either: /// -/// References are counted with multiplicity by walking each `outer` expression -/// (as `try_collapse_projection_chain` does), so a self-duplicating expression -/// such as `r + r` counts as two references. A volatile expression referenced -/// exactly once has nothing to duplicate and is left to merge. -fn would_duplicate_volatile_exprs( +/// - **volatile** (e.g. `random()`) — evaluating it independently at each site +/// makes references that should share one "locked-in" value diverge (the +/// correctness guard the physical `ProjectionPushdown` and `FilterPushdown` +/// rules also apply via +/// `datafusion_physical_expr_common::physical_expr::is_volatile`); or +/// - **not cheap to recompute** — its placement is not push-to-leaves +/// (`KeepInPlace`: arithmetic, casts, most scalar functions). Leaf-pushable +/// expressions (columns, `get_field`, `input_file_name`) still merge. This +/// matches `try_collapse_projection_chain`. +/// +/// References are counted with multiplicity, so `r + r` counts as two; an +/// expression referenced exactly once has nothing to duplicate. +fn would_duplicate_costly_exprs( inner: &ProjectionExprs, outer: &ProjectionExprs, ) -> bool { @@ -664,10 +670,10 @@ fn would_duplicate_volatile_exprs( .expect("infallible closure should not fail"); } - ref_counts - .iter() - .enumerate() - .any(|(idx, &count)| count > 1 && is_volatile(&inner_exprs[idx].expr)) + ref_counts.iter().enumerate().any(|(idx, &count)| { + let expr = &inner_exprs[idx].expr; + count > 1 && (is_volatile(expr) || !expr.placement().should_push_to_leaves()) + }) } impl DataSource for FileScanConfig { @@ -957,11 +963,13 @@ impl DataSource for FileScanConfig { projection: &ProjectionExprs, ) -> Result>> { // Don't merge a projection into the scan if it would inline a volatile - // expression that the outer projection references, which would turn a - // single "locked-in" value (e.g. `random()` aliased in a subquery) into - // multiple independent evaluations. See #23220. + // or expensive expression referenced more than once. For a volatile + // expression (e.g. `random()` aliased in a subquery) this would turn a + // single "locked-in" value into multiple independent evaluations (see + // #23220); for an expensive scalar function it would undo CSE and + // re-evaluate the expression at every reference site. if let Some(inner) = self.file_source.projection() - && would_duplicate_volatile_exprs(inner, projection) + && would_duplicate_costly_exprs(inner, projection) { return Ok(None); } @@ -3395,6 +3403,45 @@ mod tests { )) } + /// Helper: create a deterministic but expensive scalar-function + /// expression, e.g. `abs()`. + fn make_udf_expr(args: Vec>) -> Arc { + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarUDF; + use datafusion_functions::math::abs::AbsFunc; + use datafusion_physical_expr::ScalarFunctionExpr; + + Arc::new(ScalarFunctionExpr::new( + "abs", + Arc::new(ScalarUDF::from(AbsFunc::new())), + args, + Arc::new(Field::new("abs", DataType::Int32, false)), + Arc::new(ConfigOptions::default()), + )) + } + + /// Helper: create a cheap, leaf-pushable scalar function — struct field + /// access `get_field(s, 'x')`, whose placement is `MoveTowardsLeafNodes` + /// when the base is a column and the key is a literal. + fn make_leaf_pushable_expr() -> Arc { + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarUDF; + use datafusion_functions::core::getfield::GetFieldFunc; + use datafusion_physical_expr::ScalarFunctionExpr; + use datafusion_physical_expr::expressions::Literal; + + Arc::new(ScalarFunctionExpr::new( + "get_field", + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(Column::new("s", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))), + ], + Arc::new(Field::new("x", DataType::Int32, true)), + Arc::new(ConfigOptions::default()), + )) + } + /// Column-only inner projections always merge safely, even when /// the outer projection references them multiple times. #[test] @@ -3411,16 +3458,16 @@ mod tests { (Arc::new(Column::new("a", 0)), "y"), ]); - assert!(!would_duplicate_volatile_exprs(&inner, &outer)); + assert!(!would_duplicate_costly_exprs(&inner, &outer)); } - /// Deterministic computed expressions (arithmetic) referenced multiple - /// times are allowed to merge — only volatile expressions are protected. + /// A non-trivial computed expression (arithmetic, `KeepInPlace`) referenced + /// multiple times blocks the merge — recomputing it per site is wasteful. #[test] - fn test_would_duplicate_allows_deterministic_computed_multi_ref() { + fn test_would_duplicate_blocks_computed_multi_ref() { let col_a: Arc = Arc::new(Column::new("a", 0)); let col_b: Arc = Arc::new(Column::new("b", 1)); - // Inner: [a + b, b] (index 0 is deterministic computed) + // Inner: [a + b, b] (index 0 is a non-trivial computed expression) let inner = make_projection(vec![ ( Arc::new(BinaryExpr::new( @@ -3439,8 +3486,7 @@ mod tests { (Arc::new(Column::new("sum", 0)), "y"), ]); - // Deterministic arithmetic → allow merge even though duplicated - assert!(!would_duplicate_volatile_exprs(&inner, &outer)); + assert!(would_duplicate_costly_exprs(&inner, &outer)); } /// A volatile expression the outer projection does not reference is @@ -3455,7 +3501,7 @@ mod tests { // Outer references only index 1 (the column), not the volatile expr let outer = make_projection(vec![(Arc::new(Column::new("a", 1)), "a")]); - assert!(!would_duplicate_volatile_exprs(&inner, &outer)); + assert!(!would_duplicate_costly_exprs(&inner, &outer)); } /// A volatile expression referenced multiple times must block merge: @@ -3472,7 +3518,7 @@ mod tests { (Arc::new(Column::new("r", 0)), "y"), ]); - assert!(would_duplicate_volatile_exprs(&inner, &outer)); + assert!(would_duplicate_costly_exprs(&inner, &outer)); } /// A volatile expression referenced exactly once has nothing to duplicate, @@ -3490,7 +3536,7 @@ mod tests { (Arc::new(Column::new("a", 1)), "a"), ]); - assert!(!would_duplicate_volatile_exprs(&inner, &outer)); + assert!(!would_duplicate_costly_exprs(&inner, &outer)); } /// References are counted with multiplicity, so a single outer expression @@ -3510,7 +3556,7 @@ mod tests { "x", )]); - assert!(would_duplicate_volatile_exprs(&inner, &outer)); + assert!(would_duplicate_costly_exprs(&inner, &outer)); } /// A volatile expression buried inside a larger expression (e.g. @@ -3533,7 +3579,7 @@ mod tests { (Arc::new(Column::new("expr", 0)), "y"), ]); - assert!(would_duplicate_volatile_exprs(&inner, &outer)); + assert!(would_duplicate_costly_exprs(&inner, &outer)); } /// Empty projections should not block merging. @@ -3541,6 +3587,61 @@ mod tests { fn test_would_duplicate_empty_projections() { let inner = make_projection(vec![]); let outer = make_projection(vec![]); - assert!(!would_duplicate_volatile_exprs(&inner, &outer)); + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// An expensive (scalar-function) expression referenced more than once + /// must block the merge to preserve CSE. + #[test] + fn test_would_duplicate_blocks_multi_ref_expensive() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + // Inner: [abs(a)] + let inner = make_projection(vec![(make_udf_expr(vec![col_a]), "abs_a")]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("abs_a", 0)), "x"), + (Arc::new(Column::new("abs_a", 0)), "y"), + ]); + + assert!(would_duplicate_costly_exprs(&inner, &outer)); + } + + /// An expensive expression referenced only once has nothing to duplicate, + /// so the merge is allowed. + #[test] + fn test_would_duplicate_allows_single_ref_expensive() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + // Inner: [abs(a), a] + let inner = make_projection(vec![ + (make_udf_expr(vec![Arc::clone(&col_a)]), "abs_a"), + (Arc::clone(&col_a), "a"), + ]); + + // Outer references each inner column once + let outer = make_projection(vec![ + (Arc::new(Column::new("abs_a", 0)), "out"), + (Arc::new(Column::new("a", 1)), "a"), + ]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A cheap, leaf-pushable scalar function (placement + /// `MoveTowardsLeafNodes`, e.g. `get_field` / `input_file_name`) still + /// merges even when referenced multiple times — it is meant to be pushed + /// into the scan, so blocking would defeat that optimization. + #[test] + fn test_would_duplicate_allows_leaf_pushable_scalar_function() { + // Inner: [input_file_name()] + let inner = make_projection(vec![(make_leaf_pushable_expr(), "f")]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("f", 0)), "x"), + (Arc::new(Column::new("f", 0)), "y"), + ]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); } } diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index e1edca260e09f..cbbd9b74dfc00 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -6085,7 +6085,8 @@ physical_plan 03)----BoundedWindowAggExec: wdw=[sum(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, sum(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, count(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "count(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable List(Int64) }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable List(Int64) }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 04)------SortPreservingMergeExec: [c1@2 ASC NULLS LAST, c2@3 ASC NULLS LAST], fetch=5 05)--------SortExec: TopK(fetch=5), expr=[c1@2 ASC NULLS LAST, c2@3 ASC NULLS LAST], preserve_partitioning=[true] -06)----------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-0.csv, WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-1.csv], [WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-2.csv, WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-3.csv]]}, projection=[c2@1 >= 2 as __common_expr_1, c2@1 >= 2 AND c2@1 < 4 AND c1@0 > 0 as __common_expr_2, c1, c2], file_type=csv, has_header=false +06)----------ProjectionExec: expr=[__common_expr_3@0 as __common_expr_1, __common_expr_3@0 AND c2@2 < 4 AND c1@1 > 0 as __common_expr_2, c1@1 as c1, c2@2 as c2] +07)------------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-0.csv, WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-1.csv], [WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-2.csv, WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-3.csv]]}, projection=[c2@1 >= 2 as __common_expr_3, c1, c2], file_type=csv, has_header=false # FILTER filters out some rows query IIIII??