diff --git a/Cargo.lock b/Cargo.lock index 62c04d332b98e..f045e7697f576 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3284,9 +3284,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.13" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 3c1e7b50780a5..6256b0f650484 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -57,7 +57,6 @@ use crate::physical_plan::{ use crate::schema_equivalence::schema_satisfied_by; use arrow::array::{RecordBatch, builder::StringBuilder}; -use arrow::compute::SortOptions; use arrow::datatypes::Schema; use arrow_schema::Field; use datafusion_catalog::ScanArgs; @@ -1708,15 +1707,27 @@ impl DefaultPhysicalPlanner { // implement null-aware anti-join semantics and would return wrong // results when the right side contains a null join key. { - // Use SortMergeJoin if hash join is not preferred - let join_on_len = join_on.len(); + let sort_options = join_on + .iter() + .map(|(left_col, _)| { + physical_left + .output_ordering() + .and_then(|ordering| { + ordering + .iter() + .find(|sort_expr| sort_expr.expr.eq(left_col)) + .map(|sort_expr| sort_expr.options) + }) + .unwrap_or_default() + }) + .collect(); Arc::new(SortMergeJoinExec::try_new( physical_left, physical_right, join_on, join_filter, *join_type, - vec![SortOptions::default(); join_on_len], + sort_options, *null_equality, )?) } else if session_state.config().target_partitions() > 1 @@ -3316,8 +3327,8 @@ mod tests { use datafusion_expr::{ Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, HigherOrderUDF, LogicalPlanBuilder, Partitioning as LogicalPartitioning, RangePartitioning, - ScalarUDF, Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, - WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery, + ScalarUDF, Signature, SortExpr, TableSource, UserDefinedLogicalNodeCore, + Volatility, WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; @@ -3820,6 +3831,134 @@ mod tests { ctx.sql(query).await?.collect().await } + fn smj_test_context( + left_ordering: Option>, + right_ordering: Option>, + ) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![3, 2, 1])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?; + + let table = |batch, ordering| -> Result { + let table = MemTable::try_new(Arc::clone(&schema), vec![vec![batch]])?; + Ok(match ordering { + Some(ordering) => table.with_sort_order(vec![ordering]), + None => table, + }) + }; + let left = table(batch.clone(), left_ordering)?; + let right = table(batch, right_ordering)?; + + let config = SessionConfig::new() + .with_target_partitions(4) + .set_bool("datafusion.optimizer.prefer_hash_join", false) + .set_bool("datafusion.optimizer.skip_failed_rules", false); + let state = SessionStateBuilder::new() + .with_config(config) + .with_runtime_env(Arc::new(RuntimeEnv::default())) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + ctx.register_table("left_table", Arc::new(left))?; + ctx.register_table("right_table", Arc::new(right))?; + Ok(ctx) + } + + async fn smj_plan(ctx: &SessionContext, sql: &str) -> Result { + let state = ctx.state(); + let logical_plan = state.create_logical_plan(sql).await?; + let logical_plan = state.optimize(&logical_plan)?; + let plan = DefaultPhysicalPlanner::default() + .create_physical_plan(&logical_plan, &state) + .await?; + Ok(format!("{}", displayable(plan.as_ref()).indent(false))) + } + + #[tokio::test] + async fn smj_uses_left_descending_sort_options() -> Result<()> { + let descending = vec![col("a").sort(false, false)]; + let ctx = smj_test_context(Some(descending), None)?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)] + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + Ok(()) + } + + #[tokio::test] + async fn smj_defaults_to_ascending_when_left_is_unsorted() -> Result<()> { + let ctx = smj_test_context(None, None)?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + Ok(()) + } + + #[tokio::test] + async fn smj_defaults_only_unmatched_join_columns_to_ascending() -> Result<()> { + let descending = vec![col("a").sort(false, false)]; + let ctx = smj_test_context(Some(descending), None)?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table \ + ON left_table.a = right_table.a AND left_table.b = right_table.b", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0), (b@1, b@1)] + SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + SortExec: expr=[a@0 DESC NULLS LAST, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + Ok(()) + } + + #[tokio::test] + async fn smj_reuses_matching_descending_order_on_both_sides() -> Result<()> { + let descending = vec![col("a").sort(false, false)]; + let ctx = smj_test_context(Some(descending.clone()), Some(descending))?; + let plan = smj_plan( + &ctx, + "SELECT * FROM left_table JOIN right_table ON left_table.a = right_table.a", + ) + .await?; + + insta::assert_snapshot!(plan, @r" + SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)] + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + DataSourceExec: partitions=1, partition_sizes=[1], output_ordering=a@0 DESC NULLS LAST + "); + Ok(()) + } + #[tokio::test] async fn test_all_operators() -> Result<()> { let logical_plan = test_csv_scan() diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 693e9bd2cbf31..8e29a528d5b37 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -1749,7 +1749,11 @@ mod test { CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetRowSelection, RowGroupAccess, }; - use arrow::array::{RecordBatch, record_batch}; + use arrow::array::{ + Array, ArrayRef, Date32Array, Int32Array, ListArray, RecordBatch, StringArray, + record_batch, + }; + use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use bytes::{BufMut, BytesMut}; use datafusion_common::{ @@ -1774,8 +1778,8 @@ mod test { }; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_pruning::MAX_IN_LIST_SIZE; - use futures::StreamExt; use futures::stream::BoxStream; + use futures::{StreamExt, TryStreamExt}; use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}; use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData}; @@ -2340,13 +2344,12 @@ mod test { async fn collect_int32_values( mut stream: BoxStream<'static, Result>, ) -> Vec { - use arrow::array::Array; let mut values = vec![]; while let Some(Ok(batch)) = stream.next().await { let array = batch .column(0) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); for i in 0..array.len() { if !array.is_null(i) { @@ -2365,6 +2368,133 @@ mod test { write_parquet_batches(store, filename, vec![batch], None).await } + async fn read_batches_with_schema( + store: Arc, + filename: &str, + data_size: usize, + schema: SchemaRef, + ) -> Result> { + let projection_indices = (0..schema.fields().len()).collect::>(); + let file = + PartitionedFile::new(filename.to_string(), u64::try_from(data_size).unwrap()); + let morselizer = ParquetMorselizerBuilder::new() + .with_store(store) + .with_schema(schema) + .with_projection_indices(&projection_indices) + .build(); + open_file(&morselizer, file).await?.try_collect().await + } + + #[tokio::test] + async fn schema_adaptation_casts_utf8_to_date32() -> Result<()> { + let store = Arc::new(InMemory::new()) as Arc; + let file_schema = + Arc::new(Schema::new(vec![Field::new("date", DataType::Utf8, true)])); + let batch = RecordBatch::try_new( + file_schema, + vec![Arc::new(StringArray::from(vec![ + "2026-01-01", + "2026-02-01", + ]))], + )?; + let data_size = + write_parquet(Arc::clone(&store), "schema_adapt_dates.parquet", batch).await; + + let logical_schema = Arc::new(Schema::new(vec![Field::new( + "date", + DataType::Date32, + true, + )])); + let batches = read_batches_with_schema( + store, + "schema_adapt_dates.parquet", + data_size, + logical_schema, + ) + .await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + let dates = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(dates.values(), &[20_454, 20_485]); + Ok(()) + } + + #[tokio::test] + async fn schema_adaptation_normalizes_list_fields() -> Result<()> { + let store = Arc::new(InMemory::new()) as Arc; + let inner_field = Arc::new(Field::new("conditions", DataType::Int32, false)); + let file_schema = Arc::new(Schema::new(vec![Field::new( + "conditions", + DataType::List(Arc::clone(&inner_field)), + true, + )])); + let values = Int32Array::from(vec![1, 2, 3, 4]); + let offsets = OffsetBuffer::from_lengths([2, 2]); + let list = ListArray::new(inner_field, offsets, Arc::new(values), None); + let batch = RecordBatch::try_new(file_schema, vec![Arc::new(list) as ArrayRef])?; + let data_size = + write_parquet(Arc::clone(&store), "schema_adapt_list.parquet", batch).await; + + let logical_inner = Arc::new(Field::new("element", DataType::Int32, true)); + let logical_schema = Arc::new(Schema::new(vec![Field::new( + "conditions", + DataType::List(logical_inner), + true, + )])); + let batches = read_batches_with_schema( + store, + "schema_adapt_list.parquet", + data_size, + Arc::clone(&logical_schema), + ) + .await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 2); + assert_eq!(batches[0].schema(), logical_schema); + Ok(()) + } + + #[tokio::test] + async fn schema_adaptation_preserves_logical_nullability() -> Result<()> { + let store = Arc::new(InMemory::new()) as Arc; + let file_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + file_schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + let data_size = + write_parquet(Arc::clone(&store), "schema_adapt_nullable.parquet", batch) + .await; + + let logical_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, true)])); + let batches = read_batches_with_schema( + store, + "schema_adapt_nullable.parquet", + data_size, + logical_schema, + ) + .await?; + + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + assert!(batches[0].schema().field(0).is_nullable()); + let ids = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ids.values(), &[1, 2, 3]); + Ok(()) + } + /// Write multiple batches to a parquet file with optional writer properties async fn write_parquet_batches( store: Arc, @@ -3742,7 +3872,7 @@ mod test { let part = batch .column(1) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); assert!(part.iter().all(|v| v == Some(5))); diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 097b4563af5df..3d4ec2cb25081 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -507,13 +507,28 @@ impl ParquetSource { } } - #[cfg(test)] - pub(crate) fn with_reverse_row_groups(mut self, reverse_row_groups: bool) -> Self { + /// Sets whether row groups are read in reverse order. + /// + /// This changes row-group traversal only; rows within each row group retain + /// their file order. + /// + /// ``` + /// use std::sync::Arc; + /// + /// use arrow_schema::Schema; + /// use datafusion_datasource_parquet::source::ParquetSource; + /// + /// let source = ParquetSource::new(Arc::new(Schema::empty())) + /// .with_reverse_row_groups(true); + /// assert!(source.reverse_row_groups()); + /// ``` + pub fn with_reverse_row_groups(mut self, reverse_row_groups: bool) -> Self { self.reverse_row_groups = reverse_row_groups; self } - #[cfg(test)] - pub(crate) fn reverse_row_groups(&self) -> bool { + + /// Returns whether row groups are read in reverse order. + pub fn reverse_row_groups(&self) -> bool { self.reverse_row_groups } } diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index c4c1d743b58b6..fd3395cf5ba82 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -1034,3 +1034,68 @@ impl LogicalPlan { }) } } + +#[cfg(test)] +mod tests { + use std::fmt::Formatter; + use std::sync::{Arc, LazyLock}; + + use datafusion_common::{DFSchema, DFSchemaRef, Result}; + + use super::*; + use crate::UserDefinedLogicalNodeCore; + + static EMPTY_SCHEMA: LazyLock = + LazyLock::new(|| Arc::new(DFSchema::empty())); + + #[derive(Debug, Eq, Hash, PartialEq, PartialOrd)] + struct EmptyExpressionNode; + + impl UserDefinedLogicalNodeCore for EmptyExpressionNode { + fn name(&self) -> &str { + "EmptyExpressionNode" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &EMPTY_SCHEMA + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "EmptyExpressionNode") + } + + fn with_exprs_and_inputs( + &self, + exprs: Vec, + inputs: Vec, + ) -> Result { + assert!(exprs.is_empty()); + assert!(inputs.is_empty()); + Ok(Self) + } + } + + #[test] + fn map_expressions_does_not_rebuild_expressionless_extension() -> Result<()> { + let original: Arc = Arc::new(EmptyExpressionNode); + let plan = LogicalPlan::Extension(Extension { + node: Arc::clone(&original), + }); + + let transformed = plan.map_expressions(|expr| Ok(Transformed::no(expr)))?; + let LogicalPlan::Extension(Extension { node }) = transformed.data else { + panic!("expected extension plan") + }; + + assert!(Arc::ptr_eq(&node, &original)); + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index 1f9a6a583cc44..a9593d3dfd6ce 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -310,6 +310,33 @@ pub struct EquivalenceGroup { } impl EquivalenceGroup { + /// A cheap, deliberately conservative check that two groups hold the same + /// equivalence classes. + /// + /// This is not `PartialEq`, and the distinction is the point. `classes` is a + /// `Vec` whose order carries no meaning -- `remove_class_at_idx` uses + /// `swap_remove` -- so two groups describing exactly the same equalities can + /// hold their classes in different orders and this returns `false` for them. + /// Naming it `PartialEq` would invite callers to read it as semantic equality, + /// which it is not. + /// + /// The comparison is positional because it runs on a hot path: + /// `ProjectionExec` consults it every time a rule replaces its child. Set + /// semantics would mean scanning the other group once per class, and that + /// quadratic term costs more than the recomputation the caller is trying to + /// skip, by a margin that widens with the number of classes. + /// + /// Only the false direction is reachable: a group can be reported different + /// when it is not, never the same when it is not. Callers using this to skip + /// work must be built so that a `false` merely costs them that work, which is + /// exactly how the projection fast path uses it. + /// + /// `map` is an index into `classes` and carries no information the classes do + /// not already have, so it takes no part in the comparison. + pub fn has_same_classes(&self, other: &Self) -> bool { + self.classes == other.classes + } + /// Creates an equivalence group from the given equivalence classes. pub fn new(classes: impl IntoIterator) -> Self { classes.into_iter().collect::>().into() @@ -924,7 +951,7 @@ mod tests { use super::*; use crate::equivalence::tests::create_test_params; use crate::expressions::{BinaryExpr, Column, binary, col, lit}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_expr::Operator; @@ -1245,4 +1272,92 @@ mod tests { Ok(()) } + + /// Builds a group from a list of equated column pairs. + fn group_of(schema: &SchemaRef, pairs: &[(&str, &str)]) -> Result { + let mut group = EquivalenceGroup::default(); + for (lhs, rhs) in pairs { + group.add_equal_conditions(col(lhs, schema)?, col(rhs, schema)?); + } + Ok(group) + } + + fn abc_schema() -> SchemaRef { + 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), + ])) + } + + #[test] + fn test_has_same_classes_is_conservative_about_class_order() -> Result<()> { + // `classes` is a `Vec` and `remove_class_at_idx` uses `swap_remove`, so + // the order two groups hold their classes in depends on how they were + // built. This check is positional and reports such a pair as different. + // + // That is deliberate, and it is why this is not `PartialEq`: set + // semantics would mean scanning the other group per class, and the + // quadratic cost dwarfs the recomputation the caller is trying to skip. + // The error can only go this way -- different when they match, never the + // reverse -- so a caller only forfeits an optimization. + let schema = abc_schema(); + let ab_then_cd = group_of(&schema, &[("a", "b"), ("c", "d")])?; + let cd_then_ab = group_of(&schema, &[("c", "d"), ("a", "b")])?; + + assert_eq!(ab_then_cd.len(), 2, "expected two disjoint classes"); + assert!(!ab_then_cd.has_same_classes(&cd_then_ab)); + + Ok(()) + } + + #[test] + fn test_has_same_classes_compares_classes() -> Result<()> { + let schema = abc_schema(); + + // Two empty groups agree. + assert!(group_of(&schema, &[])?.has_same_classes(&group_of(&schema, &[])?)); + // The same class, built the same way. + assert!( + group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "b")])?) + ); + // A class is a set, so the order within a pair is immaterial. + assert!( + group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("b", "a")])?) + ); + // A populated group is not an empty one. + assert!( + !group_of(&schema, &[("a", "b")])?.has_same_classes(&group_of(&schema, &[])?) + ); + // Equating a different pair yields a different group. + assert!( + !group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "c")])?) + ); + // Widening a class yields a different group. + assert!( + !group_of(&schema, &[("a", "b")])? + .has_same_classes(&group_of(&schema, &[("a", "b"), ("b", "c")])?) + ); + + Ok(()) + } + + #[test] + fn test_has_same_classes_ignores_the_map() -> Result<()> { + // `map` indexes into `classes`, so equal classes must imply equal + // groups no matter how the classes were arrived at. Bridging `a = b` + // and `b = c` into one class must match stating `a = c` and `a = b`. + let schema = abc_schema(); + let bridged = group_of(&schema, &[("a", "b"), ("b", "c")])?; + let direct = group_of(&schema, &[("a", "c"), ("a", "b")])?; + + assert_eq!(bridged.len(), 1, "expected a single bridged class"); + assert!(bridged.has_same_classes(&direct)); + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 22b3382f50638..6a0fc76afe778 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -194,6 +194,31 @@ impl OrderingEquivalenceCache { } } +/// Counts how often [`EquivalenceProperties::project_reusing`] reused a cached +/// equivalence group instead of reprojecting it. +/// +/// Thread local rather than a global counter: the test binary runs tests in +/// parallel, and a shared count would let one test observe another's hits. +#[cfg(test)] +mod eq_group_reuse_probe { + use std::cell::Cell; + + thread_local! { + static HITS: Cell = const { Cell::new(0) }; + } + + pub(super) fn record_hit() { + HITS.with(|h| h.set(h.get() + 1)); + } + + /// Runs `f` and reports how many reuses happened while it did. + pub(super) fn count(f: impl FnOnce() -> T) -> (T, usize) { + let before = HITS.with(Cell::get); + let value = f(); + (value, HITS.with(Cell::get) - before) + } +} + impl EquivalenceProperties { /// Helper used by the ordering equivalence rule when considering whether /// an expression can replace an existing sort key without invalidating @@ -1169,6 +1194,48 @@ impl EquivalenceProperties { /// `output_schema`. pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self { let eq_group = self.eq_group.project(mapping); + // Built here, so it satisfies the precondition by construction; going + // through the checked entry point would reproject it under + // `debug_assertions` for nothing. + self.project_with_eq_group_unchecked(mapping, output_schema, eq_group) + } + + /// Projects `self`, reusing `cached`'s already-projected equivalence group + /// when `self`'s group is unchanged from `previous`'s. + /// + /// [`EquivalenceGroup::project`] is a pure function of the group and the + /// mapping, so when the group is unchanged the previous result can be handed + /// back rather than recomputed. Orderings are still derived here: they are + /// precisely what changes when a sort is introduced below this node. + /// + /// Falls back to a full projection when the groups differ, so there is no + /// precondition to violate. The caller does still have to pass a `cached` + /// that came from projecting `previous` through this same `mapping`; that + /// part cannot be checked here, and in practice it holds because the caller + /// carries its projection over untouched. + pub fn project_reusing( + &self, + mapping: &ProjectionMapping, + output_schema: SchemaRef, + previous: &EquivalenceProperties, + cached: &EquivalenceProperties, + ) -> Self { + let eq_group = if self.eq_group.has_same_classes(&previous.eq_group) { + #[cfg(test)] + eq_group_reuse_probe::record_hit(); + cached.eq_group.clone() + } else { + self.eq_group.project(mapping) + }; + self.project_with_eq_group_unchecked(mapping, output_schema, eq_group) + } + + fn project_with_eq_group_unchecked( + &self, + mapping: &ProjectionMapping, + output_schema: SchemaRef, + eq_group: EquivalenceGroup, + ) -> Self { let orderings = self.projected_orderings(mapping, self.oeq_cache.normal_cls.clone()); let normal_orderings = orderings @@ -1540,3 +1607,127 @@ fn get_expr_properties( expr.get_properties(&child_states) } } + +#[cfg(test)] +mod tests { + use super::*; + + use crate::equivalence::tests::create_test_params; + use crate::expressions::col; + + use arrow::datatypes::{DataType, Field, Schema}; + + /// Renames `a`, `b`, `c` and `d` so the projection is non-trivial while + /// still carrying every ordering and the `a = c` class across. + fn renaming_mapping( + schema: &SchemaRef, + output_schema: &SchemaRef, + ) -> Result { + [ + ("a", "a1", 0), + ("b", "b1", 1), + ("c", "c1", 2), + ("d", "d1", 3), + ] + .into_iter() + .map(|(source, target, index)| { + Ok(( + col(source, schema)?, + vec![(col(target, output_schema)?, index)].into(), + )) + }) + .collect::>>() + .map(|entries| entries.into_iter().collect()) + } + + fn renamed_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, true), + Field::new("b1", DataType::Int32, true), + Field::new("c1", DataType::Int32, true), + Field::new("d1", DataType::Int32, true), + ])) + } + + #[test] + fn test_project_reusing_matches_project() -> Result<()> { + let (schema, eq_properties) = create_test_params()?; + let output_schema = renamed_schema(); + let mapping = renaming_mapping(&schema, &output_schema)?; + + let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema)); + + // Reusing the projection of an identical group must reproduce it exactly. + let (reused, hits) = eq_group_reuse_probe::count(|| { + eq_properties.project_reusing( + &mapping, + Arc::clone(&output_schema), + &eq_properties, + &baseline, + ) + }); + assert_eq!( + hits, 1, + "the reuse path was not taken, so this compares nothing" + ); + + // Guard against a vacuous comparison. + assert!( + !baseline.eq_group().is_empty(), + "the projection dropped the equivalence class, nothing is being tested" + ); + assert!( + !baseline.oeq_class().is_empty(), + "the projection dropped every ordering, nothing is being tested" + ); + + assert!( + baseline.eq_group().has_same_classes(reused.eq_group()), + "equivalence group" + ); + assert_eq!(baseline.oeq_class(), reused.oeq_class(), "orderings"); + assert_eq!(baseline.constraints(), reused.constraints(), "constraints"); + assert_eq!(baseline.schema(), reused.schema(), "schema"); + + Ok(()) + } + + #[test] + fn test_project_reusing_falls_back_when_the_group_moved() -> Result<()> { + // A `previous` whose group differs must not have its projection carried + // over. There is no precondition to violate here: the fallback is what + // keeps a mismatched pair from producing properties that disagree with + // the projection. + let (schema, eq_properties) = create_test_params()?; + let output_schema = renamed_schema(); + let mapping = renaming_mapping(&schema, &output_schema)?; + + let baseline = eq_properties.project(&mapping, Arc::clone(&output_schema)); + let unrelated = EquivalenceProperties::new(Arc::clone(&schema)); + assert!( + !eq_properties + .eq_group() + .has_same_classes(unrelated.eq_group()), + "the two groups were meant to differ" + ); + + let (recomputed, hits) = eq_group_reuse_probe::count(|| { + eq_properties.project_reusing( + &mapping, + Arc::clone(&output_schema), + &unrelated, + &baseline, + ) + }); + assert_eq!(hits, 0, "a mismatched group was carried over anyway"); + + // Falling back must land on exactly what `project` would have produced. + assert!( + baseline.eq_group().has_same_classes(recomputed.eq_group()), + "equivalence group" + ); + assert_eq!(baseline.oeq_class(), recomputed.oeq_class(), "orderings"); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index adc8f8c315b32..5f4e6a70b37be 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -106,6 +106,54 @@ pub trait ArrowHashTable { fn null_map_idxs(&self) -> Vec; } +enum StringArrayType { + Utf8(StringArray), + Utf8View(StringViewArray), + LargeUtf8(LargeStringArray), +} + +impl StringArrayType { + fn value(&self, row_idx: usize) -> Option<&str> { + match self { + Self::Utf8(array) if !array.is_null(row_idx) => Some(array.value(row_idx)), + Self::Utf8View(array) if !array.is_null(row_idx) => { + Some(array.value(row_idx)) + } + Self::LargeUtf8(array) if !array.is_null(row_idx) => { + Some(array.value(row_idx)) + } + _ => None, + } + } +} + +impl TryFrom<&DataType> for StringArrayType { + type Error = DataType; + + fn try_from(data_type: &DataType) -> std::result::Result { + let values = Vec::<&str>::new(); + match data_type { + DataType::Utf8 => Ok(Self::Utf8(values.into())), + DataType::Utf8View => Ok(Self::Utf8View(values.into())), + DataType::LargeUtf8 => Ok(Self::LargeUtf8(values.into())), + data_type => Err(data_type.clone()), + } + } +} + +impl TryFrom for StringArrayType { + type Error = DataType; + + fn try_from(array: ArrayRef) -> std::result::Result { + match array.data_type() { + DataType::Utf8 => Ok(Self::Utf8(array.as_string::().clone())), + DataType::Utf8View => Ok(Self::Utf8View(array.as_string_view().clone())), + DataType::LargeUtf8 => Ok(Self::LargeUtf8(array.as_string::().clone())), + data_type => Err(data_type.clone()), + } + } +} + /// 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) @@ -121,10 +169,10 @@ pub fn is_supported_hash_key_type(kt: &DataType) -> bool { // An implementation of ArrowHashTable for String keys pub struct StringHashTable { - owned: ArrayRef, + data_type: DataType, + owned: StringArrayType, map: TopKHashTable>, rnd: RandomState, - data_type: DataType, } // An implementation of ArrowHashTable for any `ArrowPrimitiveType` key @@ -132,69 +180,37 @@ struct PrimitiveHashTable where Option<::Native>: Comparable, { - owned: ArrayRef, + data_type: DataType, + owned: PrimitiveArray, map: TopKHashTable>, 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"), + pub fn new(limit: usize, data_type: &DataType) -> Self { + let Ok(owned) = StringArrayType::try_from(data_type) else { + unreachable!("StringHashTable requires a string data type") }; - Self { + data_type: data_type.clone(), 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 { - let is_null_and_value = match self.data_type { - DataType::Utf8 => { - let arr = self.owned.as_string::(); - (arr.is_null(row_idx), arr.value(row_idx)) - } - DataType::LargeUtf8 => { - let arr = self.owned.as_string::(); - (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, 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; + assert_eq!( + ids.data_type(), + &self.data_type, + "Top-K hash key batch type must match the declared type" + ); + let Ok(owned) = StringArrayType::try_from(ids) else { + unreachable!("StringHashTable requires a string array") + }; + self.owned = owned; } fn len(&self) -> usize { @@ -211,11 +227,11 @@ impl ArrowHashTable for StringHashTable { fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); - match self.data_type { + 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!(), + DataType::LargeUtf8 => Arc::new(LargeStringArray::from(ids)), + _ => unreachable!("StringHashTable requires a string data type"), } } @@ -224,27 +240,26 @@ impl ArrowHashTable for StringHashTable { 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| id_for_eq.as_deref() == mi.as_deref(); + let id = self.owned.value(row_idx); + let hash = self.rnd.hash_one(id); + let eq = move |stored: &Option| id == stored.as_deref(); // Use entry API to avoid double lookup - self.map.find_or_insert(hash, id, replace_idx, eq) + self.map + .find_or_insert(hash, id.map(ToOwned::to_owned), 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| id_for_eq.as_deref() == mi.as_deref(); - self.map.insert_null(hash, id, eq) + let id = self.owned.value(row_idx); + let hash = self.rnd.hash_one(id); + let eq = move |stored: &Option| id == stored.as_deref(); + self.map.insert_null(hash, id.map(ToOwned::to_owned), 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| id.as_deref() == mi.as_deref(); + let id = self.owned.value(row_idx); + let hash = self.rnd.hash_one(id); + let eq = move |stored: &Option| id == stored.as_deref(); self.map.remove_if_null(hash, eq) } @@ -255,30 +270,26 @@ impl ArrowHashTable for StringHashTable { impl PrimitiveHashTable where - Option<::Native>: Comparable, - Option<::Native>: HashValue, + Option<::Native>: Comparable + HashValue, { pub fn new(limit: usize, kt: DataType) -> Self { - let owned = Arc::new( - PrimitiveArray::::builder(0) - .with_data_type(kt.clone()) - .finish(), - ); + let owned = PrimitiveArray::::builder(0) + .with_data_type(kt.clone()) + .finish(); Self { + data_type: kt, 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, u64) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { + let id = if self.owned.is_null(row_idx) { None } else { - Some(ids.value(row_idx)) + Some(self.owned.value(row_idx)) }; let hash: u64 = id.hash(&self.rnd); (id, hash) @@ -287,11 +298,15 @@ where impl ArrowHashTable for PrimitiveHashTable where - Option<::Native>: Comparable, - Option<::Native>: HashValue, + Option<::Native>: Comparable + HashValue, { fn set_batch(&mut self, ids: ArrayRef) { - self.owned = ids; + assert_eq!( + ids.data_type(), + &self.data_type, + "Top-K hash key batch type must match the declared type" + ); + self.owned = ids.as_primitive().clone(); } fn len(&self) -> usize { @@ -309,7 +324,7 @@ where fn take_all(&mut self, indexes: Vec) -> ArrayRef { let ids = self.map.take_all(indexes); let mut builder: PrimitiveBuilder = - PrimitiveArray::builder(ids.len()).with_data_type(self.kt.clone()); + PrimitiveArray::builder(ids.len()).with_data_type(self.data_type.clone()); for id in ids.into_iter() { match id { None => builder.append_null(), @@ -325,11 +340,10 @@ where row_idx: usize, replace_idx: usize, ) -> (usize, InsertKind) { - let ids = self.owned.as_primitive::(); - let id: Option = if ids.is_null(row_idx) { + let id = if self.owned.is_null(row_idx) { None } else { - Some(ids.value(row_idx)) + Some(self.owned.value(row_idx)) }; // Compute hash and create equality closure for hash table lookup. let hash: u64 = id.hash(&self.rnd); @@ -597,9 +611,9 @@ pub fn new_hash_table( 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))), + 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))), _ => {} } @@ -629,6 +643,27 @@ mod tests { Ok(()) } + #[test] + #[should_panic(expected = "Top-K hash key batch type must match the declared type")] + fn should_reject_primitive_key_metadata_mismatch() { + let ids = TimestampMillisecondArray::from(vec![1000]); + let data_type = DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())); + let mut hash_table = + new_hash_table(1, data_type).expect("create timestamp hash table"); + + hash_table.set_batch(Arc::new(ids)); + } + + #[test] + #[should_panic(expected = "Top-K hash key batch type must match the declared type")] + fn should_reject_string_key_type_mismatch() { + let ids = LargeStringArray::from(vec!["value"]); + let mut hash_table = + new_hash_table(1, DataType::Utf8).expect("create string hash table"); + + hash_table.set_batch(Arc::new(ids)); + } + #[test] fn should_resize_properly() -> Result<()> { let mut heap_to_map = BTreeMap::::new(); diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index ca321cdf99784..3659c9757c7ae 100644 --- a/datafusion/physical-plan/src/aggregates/topk/heap.rs +++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs @@ -95,7 +95,7 @@ pub struct PrimitiveHeap where ::Native: Comparable, { - batch: ArrayRef, + batch: PrimitiveArray, heap: TopKHeap, desc: bool, data_type: DataType, @@ -106,9 +106,11 @@ where ::Native: Comparable, { pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self { - let owned: ArrayRef = Arc::new(PrimitiveArray::::builder(0).finish()); + let batch = PrimitiveArray::::builder(0) + .with_data_type(data_type.clone()) + .finish(); Self { - batch: owned, + batch, heap: TopKHeap::new(limit, desc), desc, data_type, @@ -121,15 +123,14 @@ where ::Native: Comparable, { fn set_batch(&mut self, vals: ArrayRef) { - self.batch = vals; + self.batch = vals.as_primitive().clone(); } fn is_worse(&self, row_idx: usize) -> bool { if !self.heap.is_full() { return false; } - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.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) } @@ -139,8 +140,7 @@ where } fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) { - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); self.heap.append_or_replace(new_val, map_idx, map); } @@ -150,8 +150,7 @@ where row_idx: usize, map: &mut Vec<(usize, usize)>, ) { - let vals = self.batch.as_primitive::(); - let new_val = vals.value(row_idx); + let new_val = self.batch.value(row_idx); self.heap.replace_if_better(heap_idx, new_val, map); } @@ -443,14 +442,13 @@ impl TopKHeap { } 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"); + self.heap.swap(a_idx, b_idx); - mapper.push((a_hi.map_idx, b_idx)); - mapper.push((b_hi.map_idx, a_idx)); + let b_hi = self.heap[b_idx].as_ref().expect("Missing heap entry"); + let a_hi = self.heap[a_idx].as_ref().expect("Missing heap entry"); - self.heap[a_idx] = Some(b_hi); - self.heap[b_idx] = Some(a_hi); + mapper.push((b_hi.map_idx, b_idx)); + mapper.push((a_hi.map_idx, a_idx)); } fn heapify_down(&mut self, node_idx: usize, mapper: &mut Vec<(usize, usize)>) { diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index cf362cdee55d3..1ec278eb9c377 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -52,6 +52,7 @@ use datafusion_common::tree_node::{ use datafusion_common::{DataFusionError, JoinSide, Result, internal_err, plan_err}; use datafusion_execution::TaskContext; use datafusion_expr::ExpressionPlacement; +use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::projection::Projector; use datafusion_physical_expr_common::physical_expr::{PhysicalExprRef, fmt_sql}; @@ -177,6 +178,32 @@ impl ProjectionExec { fn try_from_projector( projector: Projector, input: Arc, + ) -> Result { + Self::try_from_projector_with_eq_group(projector, input, None) + } + + /// As [`Self::try_from_projector`], but `reuse_from` may carry the previous + /// child's equivalence properties together with the projection they produced, + /// letting [`EquivalenceProperties::project_reusing`] skip reprojecting a + /// group that has not changed. + /// + /// Projecting an equivalence group is a pure function of the group and the + /// mapping, so reuse is sound exactly when both are unchanged. + /// + /// The caller establishes the first by comparing the old and new child + /// groups. The second holds because the mapping comes from + /// `projector.projection()`, carried over untouched, and from the child's + /// schema, which `ProjectionMapping::try_new` consults only for field names + /// and indices -- never for types or nullability. So a child differing only + /// in nullability keeps the same mapping. A child that renamed or reordered + /// those fields would change the group too, since its members are `Column`s + /// carrying those names, and the comparison above would reject it; were one + /// to slip through anyway, `try_new`'s name assertion errors out rather than + /// letting a stale group into the plan. + fn try_from_projector_with_eq_group( + projector: Projector, + input: Arc, + reuse_from: Option<(&EquivalenceProperties, &EquivalenceProperties)>, ) -> Result { // Construct a map from the input expressions to the output expression of the Projection let projection_mapping = @@ -185,6 +212,7 @@ impl ProjectionExec { &input, &projection_mapping, Arc::clone(projector.output_schema()), + reuse_from, )?; Ok(Self { projector, @@ -214,10 +242,21 @@ impl ProjectionExec { input: &Arc, projection_mapping: &ProjectionMapping, schema: SchemaRef, + reuse_from: Option<(&EquivalenceProperties, &EquivalenceProperties)>, ) -> Result { - // Calculate equivalence properties: + // Calculate equivalence properties. Whether the group is reprojected or + // handed back is the only thing reuse changes; everything below is + // common, so the two paths cannot drift apart. let input_eq_properties = input.equivalence_properties(); - let eq_properties = input_eq_properties.project(projection_mapping, schema); + let eq_properties = match reuse_from { + Some((previous, cached)) => input_eq_properties.project_reusing( + projection_mapping, + schema, + previous, + cached, + ), + None => input_eq_properties.project(projection_mapping, schema), + }; // Calculate output partitioning, which needs to respect aliases: let output_partitioning = input .output_partitioning() @@ -352,11 +391,28 @@ impl ExecutionPlan for ProjectionExec { metrics: ExecutionPlanMetricsSet::new(), ..Self::clone(&*self) })), - ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), - ) - .map(|p| Arc::new(p) as _), + ChildrenPropertiesMode::Recompute => { + // `Keep` above requires the child's properties to be unchanged + // outright. A rule that introduces a sort below this projection + // does not qualify, yet the child's *equivalence group* is still + // identical: sorting changes which orderings hold, not which + // expressions are equal to one another. Projecting that group + // again would reproduce the group already cached here, so reuse + // it and derive only the orderings. + // Hand over what this projection was built from and what that + // produced; `project_reusing` decides whether the group can be + // carried over and falls back to a full projection otherwise. + let reuse_from = Some(( + self.input.equivalence_properties(), + self.cache.equivalence_properties(), + )); + ProjectionExec::try_from_projector_with_eq_group( + self.projector.clone(), + children.swap_remove(0), + reuse_from, + ) + .map(|p| Arc::new(p) as _) + } } } @@ -1436,6 +1492,8 @@ mod tests { use crate::common::collect; use crate::empty::EmptyExec; + use crate::filter::FilterExec; + use crate::sorts::sort::SortExec; use crate::filter_pushdown::PushedDown; use crate::statistics::{StatisticsArgs, StatisticsContext}; @@ -2215,4 +2273,281 @@ mod tests { Ok(()) } + + /// `EmptyExec(a, b, c)` under a filter that equates `lhs` and `rhs`, so the + /// child carries a non-trivial equivalence group. + fn filtered_source(lhs: &str, rhs: &str) -> Result> { + filtered_source_with_nullability(lhs, rhs, false) + } + + /// As [`filtered_source`], but `nullable` varies the schema's nullability + /// while leaving field names and order alone. + fn filtered_source_with_nullability( + lhs: &str, + rhs: &str, + nullable: bool, + ) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, nullable), + Field::new("b", DataType::Int32, nullable), + Field::new("c", DataType::Int32, nullable), + ])); + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let predicate = binary( + col(lhs, &schema)?, + Operator::Eq, + col(rhs, &schema)?, + &schema, + )?; + Ok(Arc::new(FilterExec::try_new(predicate, input)?)) + } + + /// `[a AS x, b AS y, c AS z]` against `filtered_source`'s schema. + fn renaming_exprs(schema: &SchemaRef) -> Result> { + [("a", "x"), ("b", "y"), ("c", "z")] + .into_iter() + .map(|(source, alias)| { + Ok(ProjectionExpr { + expr: col(source, schema)?, + alias: alias.to_string(), + }) + }) + .collect() + } + + fn assert_same_properties(actual: &dyn ExecutionPlan, expected: &ProjectionExec) { + let actual_props = actual.properties().equivalence_properties(); + let expected_props = expected.properties().equivalence_properties(); + assert!( + actual_props + .eq_group() + .has_same_classes(expected_props.eq_group()), + "equivalence group: {:?} vs {:?}", + actual_props.eq_group(), + expected_props.eq_group() + ); + assert_eq!( + actual_props.oeq_class(), + expected_props.oeq_class(), + "orderings" + ); + assert_eq!( + actual_props.constraints(), + expected_props.constraints(), + "constraints" + ); + assert_eq!(actual_props.schema(), expected_props.schema(), "schema"); + // `Partitioning` has no `PartialEq`, so compare the partition count and + // the explicit `Display` form. Derived `Debug` would change with any + // field addition, making this brittle for no gain. + let actual_partitioning = actual.properties().output_partitioning(); + let expected_partitioning = expected.properties().output_partitioning(); + assert_eq!( + actual_partitioning.partition_count(), + expected_partitioning.partition_count(), + "partition count" + ); + assert_eq!( + actual_partitioning.to_string(), + expected_partitioning.to_string(), + "partitioning" + ); + } + + #[test] + fn test_sort_below_changes_orderings_but_not_the_equivalence_group() -> Result<()> { + // The premise the fast path rests on. If a sort ever starts altering + // the equivalence group, reusing the cached group becomes unsound and + // this test is the one that should fail first. + let child = filtered_source("a", "b")?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "c", + &child.schema(), + )?)]) + .expect("non-empty ordering"); + let sorted = SortExec::new(ordering, Arc::clone(&child)); + + let child_props = child.properties().equivalence_properties(); + let sorted_props = sorted.properties().equivalence_properties(); + + assert!( + !child_props.eq_group().is_empty(), + "the filter did not produce an equivalence class" + ); + assert!( + child_props + .eq_group() + .has_same_classes(sorted_props.eq_group()), + "sorting altered the equivalence group" + ); + assert_ne!( + child_props.oeq_class(), + sorted_props.oeq_class(), + "sorting did not alter the orderings" + ); + + Ok(()) + } + + #[test] + fn test_replace_children_reuses_eq_group_when_only_orderings_change() -> Result<()> { + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = + Arc::new(ProjectionExec::try_new(exprs.clone(), Arc::clone(&child))?); + + // Sorting below the projection changes which orderings hold but leaves + // the equivalence group untouched -- the case the fast path targets. + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "c", + &child.schema(), + )?)]) + .expect("non-empty ordering"); + let sorted: Arc = + Arc::new(SortExec::new(ordering, Arc::clone(&child))); + + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&sorted)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + // Guard against vacuity: the group must be worth reusing, and the sort + // must genuinely have added an ordering the original did not have. + assert!( + !projection + .properties() + .equivalence_properties() + .eq_group() + .is_empty(), + "the projection carries no equivalence class, nothing is being reused" + ); + assert!( + projection + .properties() + .equivalence_properties() + .oeq_class() + .is_empty(), + "the unsorted projection was already ordered" + ); + assert!( + !replaced + .properties() + .equivalence_properties() + .oeq_class() + .is_empty(), + "the sort did not introduce an ordering" + ); + + // The fast path must agree with building the projection from scratch. + let expected = ProjectionExec::try_new(exprs, sorted)?; + assert_same_properties(replaced.as_ref(), &expected); + + Ok(()) + } + + #[test] + fn test_replace_children_reuses_eq_group_across_a_nullability_change() -> Result<()> { + // Reuse is sound only if the projection mapping is unchanged as well. + // `ProjectionMapping::try_new` reads the child schema for field names + // and indices alone, so a child differing only in nullability keeps the + // same mapping and must still take the fast path. + // + // The swap tightens nullability rather than loosening it, matching what + // `is_allowed_field_change` permits of a physical optimizer rule. + // + // The comparison is against `try_from_projector`, the path this one + // replaces, rather than a freshly built projection: `replace_children` + // carries the existing `Projector` over, so the output schema stays as + // it was, while `try_new` would derive a new one from the new child. + // That difference is inherent to `replace_children` and not something + // this fast path introduces, so the meaningful contract is that the two + // `replace_children` paths agree. + let child = filtered_source_with_nullability("a", "b", true)?; + let exprs = renaming_exprs(&child.schema())?; + let projection = Arc::new(ProjectionExec::try_new(exprs, Arc::clone(&child))?); + + let tightened_child = filtered_source_with_nullability("a", "b", false)?; + assert_ne!( + child.schema(), + tightened_child.schema(), + "the two children were meant to differ in nullability" + ); + assert!( + child + .properties() + .equivalence_properties() + .eq_group() + .has_same_classes( + tightened_child + .properties() + .equivalence_properties() + .eq_group() + ), + "nullability moved the equivalence group, so the fast path is no longer under test" + ); + + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&tightened_child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let recomputed = ProjectionExec::try_from_projector( + projection.projector.clone(), + tightened_child, + )?; + + assert_same_properties(replaced.as_ref(), &recomputed); + + Ok(()) + } + + #[test] + fn test_replace_children_recomputes_when_eq_group_changes() -> Result<()> { + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = + Arc::new(ProjectionExec::try_new(exprs.clone(), Arc::clone(&child))?); + + // This child equates a different pair, so the cached group is stale and + // reusing it would be unsound: the guard has to fall through. + let other = filtered_source("a", "c")?; + let replaced = Arc::clone(&projection).replace_children( + vec![Arc::clone(&other)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + assert!( + !replaced + .properties() + .equivalence_properties() + .eq_group() + .has_same_classes( + projection.properties().equivalence_properties().eq_group() + ), + "the projection kept the previous child's equivalence group" + ); + + let expected = ProjectionExec::try_new(exprs, other)?; + assert_same_properties(replaced.as_ref(), &expected); + + Ok(()) + } + + #[test] + fn test_replace_children_keep_mode_carries_properties_over() -> Result<()> { + // `Keep` is the caller's promise that the new child's properties match + // the old one's, so the cached properties must survive verbatim rather + // than being derived again. + let child = filtered_source("a", "b")?; + let exprs = renaming_exprs(&child.schema())?; + let projection = Arc::new(ProjectionExec::try_new(exprs, Arc::clone(&child))?); + + let replaced = Arc::clone(&projection).replace_children( + vec![filtered_source("a", "b")?], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + )?; + + assert_same_properties(replaced.as_ref(), &projection); + + Ok(()) + } } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index c1cc5da31abaf..de85e5c20bc14 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -746,12 +746,14 @@ impl ExecutionPlan for InterleaveExec { ..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)?)) + if can_interleave(children.iter()) { + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } else { + // An optimizer can legitimately change child partitioning after + // introducing InterleaveExec. UnionExec preserves correctness when + // the rewritten children can no longer share output partitions. + UnionExec::try_new(children) + } } } } @@ -1764,6 +1766,40 @@ mod tests { Ok(()) } + #[test] + fn interleave_child_rewrite_falls_back_to_union() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let interleave: Arc = + Arc::new(InterleaveExec::try_new(vec![ + make_hash_exec(&schema, vec!["a"], 3)?, + make_hash_exec(&schema, vec!["a"], 3)?, + ])?); + + let compatible = Arc::clone(&interleave).replace_children( + vec![ + make_hash_exec(&schema, vec!["b"], 3)?, + make_hash_exec(&schema, vec!["b"], 3)?, + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + assert!(compatible.downcast_ref::().is_some()); + assert_eq!(compatible.output_partitioning().partition_count(), 3); + + let incompatible = interleave.replace_children( + vec![ + make_hash_exec(&schema, vec!["a"], 3)?, + make_hash_exec(&schema, vec!["b"], 3)?, + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + assert!(incompatible.downcast_ref::().is_some()); + assert_eq!(incompatible.output_partitioning().partition_count(), 6); + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 27d1101036d9b..748d77d63bf4c 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -717,4 +717,4 @@ enum MetricCategory { message ExplainAnalyzeCategoriesNode { bool all = 1; repeated MetricCategory only = 2; -} \ No newline at end of file +}