diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 00e016ae02cad..fbdf682f0c20c 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -15,18 +15,22 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; -use std::sync::Arc; +use std::{any::Any, fs, sync::Arc}; +use arrow::array::Float64Array; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::CsvSource; use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_plan::collect; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion::test::object_store::local_unpartitioned_file; use datafusion_common::config::{ConfigOptions, CsvOptions}; use datafusion_common::{JoinSide, JoinType, NullEquality, Result, ScalarValue}; use datafusion_datasource::TableSchema; +use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; @@ -34,9 +38,11 @@ use datafusion_expr::{ Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; use datafusion_expr_common::columnar_value::ColumnarValue; +use datafusion_functions::math::random::RandomFunc; use datafusion_physical_expr::expressions::{ BinaryExpr, CaseExpr, CastExpr, Column, Literal, NegativeExpr, binary, cast, col, }; +use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::{Distribution, Partitioning, ScalarFunctionExpr}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ @@ -63,6 +69,7 @@ use datafusion_physical_plan::{ExecutionPlan, displayable}; use insta::assert_snapshot; use itertools::Itertools; +use tempfile::NamedTempFile; /// Mocked UDF #[derive(Debug, PartialEq, Eq, Hash)] @@ -445,6 +452,63 @@ fn create_projecting_memory_exec() -> Arc { MemorySourceConfig::try_new_exec(&[], schema, Some(vec![2, 0, 3, 4])).unwrap() } +#[tokio::test] +async fn test_volatile_projection_pushdown_does_not_duplicate_evaluation() -> Result<()> { + let file = NamedTempFile::new()?; + fs::write(file.path(), "1\n2\n3\n")?; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let random = Arc::new(ScalarFunctionExpr::try_new( + Arc::new(ScalarUDF::from(RandomFunc::new())), + vec![], + schema.as_ref(), + Arc::new(ConfigOptions::default()), + )?); + let projection = ProjectionExprs::new(vec![ProjectionExpr::new(random, "r")]); + let source: Arc = + Arc::new(CsvSource::new(schema).with_csv_options(CsvOptions { + has_header: Some(false), + ..Default::default() + })); + let source = source.try_pushdown_projection(&projection)?.unwrap(); + let source = DataSourceExec::from_data_source( + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(local_unpartitioned_file(file.path()).into()) + .build(), + ); + + let outer: Arc = Arc::new(ProjectionExec::try_new( + vec![ + ProjectionExpr::new(Arc::new(Column::new("r", 0)), "left"), + ProjectionExpr::new(Arc::new(Column::new("r", 0)), "right"), + ], + source, + )?); + + let mut options = ConfigOptions::new(); + options.execution.target_partitions = 1; + let optimized = ProjectionPushdown::new().optimize(outer, &options)?; + let context = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + let batches = collect(optimized, context.task_ctx()).await?; + + for batch in batches { + let left = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let right = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(left, right); + } + + Ok(()) +} + #[test] fn test_csv_after_projection() -> Result<()> { let csv = create_projecting_csv_exec(); diff --git a/datafusion/datasource/src/file_scan_config.rs b/datafusion/datasource/src/file_scan_config.rs index c3e5cabce7bc2..a5ca8edbac8cc 100644 --- a/datafusion/datasource/src/file_scan_config.rs +++ b/datafusion/datasource/src/file_scan_config.rs @@ -41,7 +41,7 @@ use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; -use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, is_volatile}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_plan::SortOrderPushdownResult; use datafusion_physical_plan::coop::cooperative; @@ -826,6 +826,11 @@ impl DataSource for FileScanConfig { &self, projection: &ProjectionExprs, ) -> Result>> { + if let Some(inner) = self.file_source.projection() + && would_duplicate_volatile_exprs(inner, projection) + { + return Ok(None); + } match self.file_source.try_pushdown_projection(projection)? { Some(new_source) => { let mut new_file_scan_config = self.clone(); @@ -925,6 +930,36 @@ impl DataSource for FileScanConfig { } } +/// Returns `true` if merging `outer` into `inner` would duplicate a volatile +/// expression. +fn would_duplicate_volatile_exprs( + inner: &ProjectionExprs, + outer: &ProjectionExprs, +) -> bool { + use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + + let inner_exprs = inner.as_ref(); + let mut ref_counts = vec![0usize; inner_exprs.len()]; + for proj_expr in outer.as_ref() { + proj_expr + .expr + .apply(|expr| { + if let Some(col) = expr.as_any().downcast_ref::() + && let Some(count) = ref_counts.get_mut(col.index()) + { + *count += 1; + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("infallible closure should not fail"); + } + + ref_counts + .iter() + .enumerate() + .any(|(idx, &count)| count > 1 && is_volatile(&inner_exprs[idx].expr)) +} + impl FileScanConfig { /// Returns only the output orderings that are validated against actual /// file group statistics.