diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 9338fcc0bce35..d94253a84aa5f 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -33,7 +33,7 @@ use arrow::compute::{SortOptions}; use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::config::{ConfigOptions, CsvOptions}; use datafusion_common::tree_node::{TreeNode, TransformedResult}; -use datafusion_common::{create_array, NullEquality, Result, TableReference}; +use datafusion_common::{create_array, DataFusionError, NullEquality, Result, TableReference}; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::source::DataSourceExec; use datafusion_expr_common::operator::Operator; @@ -49,7 +49,7 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::sorts::sort::SortExec; -use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan}; +use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan, ExecutionPlanProperties}; use datafusion::datasource::physical_plan::CsvSource; use datafusion::datasource::listing::PartitionedFile; use datafusion_physical_optimizer::enforce_sorting::{PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting}; @@ -60,12 +60,15 @@ use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; use datafusion_physical_optimizer::output_requirements::OutputRequirementExec; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion::prelude::*; -use arrow::array::{record_batch, ArrayRef, Int32Array, RecordBatch}; +use arrow::array::{record_batch, Array, ArrayRef, Int32Array, RecordBatch}; use arrow::datatypes::{Field}; use arrow_schema::Schema; use datafusion_execution::TaskContext; use datafusion_catalog::streaming::StreamingTable; +use datafusion_expr_common::columnar_value::ColumnarValue; +use datafusion_physical_expr::projection::ProjectionExpr; +use datafusion_physical_plan::projection::ProjectionExec; use futures::StreamExt; use insta::{Settings, assert_snapshot}; @@ -3255,3 +3258,151 @@ async fn test_does_not_push_fetch_sort_through_projection_over_union() -> Result Ok(()) } + +/// A pass-through wrapper around a column: just assert that column does not contain any nulls +#[derive(Debug, Eq)] +struct AssertNotNull { + inner: Arc, +} + +impl AssertNotNull { + fn new(inner: Arc) -> Arc { + Arc::new(Self { inner }) + } +} + +impl PartialEq for AssertNotNull { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +impl std::hash::Hash for AssertNotNull { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + } +} + +impl std::fmt::Display for AssertNotNull { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "assert_not_null({})", self.inner) + } +} + +impl PhysicalExpr for AssertNotNull { + fn data_type(&self, input_schema: &Schema) -> Result { + self.inner.data_type(input_schema) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let child = self.inner.evaluate(batch)?; + match child { + ColumnarValue::Array(a) if a.logical_null_count() > 0 => Err( + DataFusionError::Internal("AssertNotNull evaluated to null".to_string()), + ), + ColumnarValue::Scalar(s) if s.is_null() => Err(DataFusionError::Internal( + "AssertNotNull evaluated to null".to_string(), + )), + child => Ok(child), + } + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(AssertNotNull { + inner: Arc::clone(&children[0]), + })) + } + + fn get_properties( + &self, + children: &[datafusion_expr::sort_properties::ExprProperties], + ) -> Result { + Ok(children[0].clone()) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "assert_not_null({})", self.inner) + } +} + +#[tokio::test] +async fn test_passthrough_wrapper_projection_keeps_ordering() -> Result<()> { + fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col(name, schema).unwrap(), + options: Default::default(), + } + } + + pub fn projection_exec( + expr: Vec<(Arc, String)>, + input: Arc, + ) -> Result> { + let proj_exprs: Vec = expr + .into_iter() + .map(|(expr, alias)| ProjectionExpr { expr, alias }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) + } + + let batch = record_batch!( + ("a", Utf8, ["x", "y"]), + ("b", Utf8, ["1", "2"]), + ("c", Utf8, ["1", "2"]) + )?; + let schema = batch.schema(); + let source = Arc::new(DataSourceExec::new(Arc::new( + datafusion::datasource::memory::MemorySourceConfig::try_new( + &[vec![batch]], + schema.clone(), + None, + )? + .try_with_sort_information(vec![ + LexOrdering::new([ + sort_expr("a", &schema), + sort_expr("b", &schema), + sort_expr("c", &schema), + ]) + .unwrap(), + ])?, + ))) as Arc; + + let projection = projection_exec( + vec![ + (AssertNotNull::new(col("a", &schema)?), "a".to_string()), + (AssertNotNull::new(col("b", &schema)?), "b".to_string()), + (AssertNotNull::new(col("c", &schema)?), "c".to_string()), + ], + source, + )?; + + let ordering = LexOrdering::new([ + sort_expr("a", &projection.schema()), + sort_expr("b", &projection.schema()), + sort_expr("c", &projection.schema()), + ]) + .unwrap(); + + let sort_satisfied = projection + .equivalence_properties() + .ordering_satisfy(ordering.clone())?; + + let plan_str = displayable(projection.as_ref()).indent(true).to_string(); + assert!( + sort_satisfied, + "sort should be satisfied, ordering: {ordering}\nplan:\n{plan_str}" + ); + + Ok(()) +} diff --git a/datafusion/expr-common/src/sort_properties.rs b/datafusion/expr-common/src/sort_properties.rs index 04da574882d30..74d644f79faef 100644 --- a/datafusion/expr-common/src/sort_properties.rs +++ b/datafusion/expr-common/src/sort_properties.rs @@ -141,7 +141,61 @@ pub struct ExprProperties { pub range: Interval, /// Indicates whether the expression preserves lexicographical ordering /// of its inputs. + /// + /// This is a *non-strict* (monotone) property: inputs advancing in + /// lexicographical order never make the output decrease, but distinct + /// inputs may map to equal outputs (ties). See + /// [`Self::strictly_order_preserving`] for the strict variant and an + /// explanation of the difference. pub preserves_lex_ordering: bool, + /// Indicates whether the expression is strictly order-preserving with + /// respect to its inputs that are `Ordered`: the output is ordered in the + /// same direction, equal outputs can only result from equal values of + /// those inputs (i.e. the mapping is one-to-one), and nulls map to nulls. + /// + /// i.e. setting this to true means that `a.cmp(b) == f(a).cmp(f(b))` + /// + /// # Difference from [`Self::preserves_lex_ordering`] + /// + /// The two properties differ in both their premise and their strictness: + /// + /// - `preserves_lex_ordering` assumes the inputs advance in + /// *lexicographical* order (a later input may decrease whenever an + /// earlier one increases), and only promises a non-decreasing output, + /// allowing distinct inputs to collapse into equal outputs; `floor`, + /// `date_trunc` and narrowing casts do exactly that. + /// - `strictly_order_preserving` assumes every `Ordered` input advances + /// *simultaneously* (component-wise, which is what actually holds when + /// all of them are sorted in the data), and promises a strict output: + /// equal outputs only from equal inputs. + /// + /// For an expression with a single ordered input the premises coincide, + /// and this field is simply the stronger claim: it implies + /// `preserves_lex_ordering`. With multiple ordered inputs, neither + /// implies the other: a lexicographical-ordering-preserving expression + /// need not be strict (distinct inputs may still produce equal outputs), + /// while `a + b` over two ordered, overflow-free inputs is strict but not + /// lexicographical (under the lexicographical premise `b` may decrease + /// while `a` increases, making the sum decrease). + /// + /// The distinction matters for suffix sort keys. Optimizers use this + /// field to substitute a sort key with an expression computed from it: + /// if data is sorted by `[x, y]`, it is also sorted by `[expr(x), y]`. + /// That claim requires `y` to be sorted within each run of equal + /// `expr(x)` values, which only holds if equal outputs imply equal `x` + /// values. With a merely monotone expression such as `floor`, one output + /// run can span several `x` groups, and `y` restarts at each group: + /// + /// ```text + /// sorted by [x, y]: (1.2, 5), (1.8, 1), (2.5, 3) + /// [floor(x), y]: (1, 5), (1, 1), (2, 3) <-- y not sorted within + /// the "1" run + /// ``` + /// + /// Hence a monotone expression only justifies the length-1 ordering + /// `[expr(x)]`, while a strictly order-preserving one keeps the entire + /// suffix valid. When in doubt, set to `false`. + pub strictly_order_preserving: bool, } impl ExprProperties { @@ -152,6 +206,7 @@ impl ExprProperties { sort_properties: SortProperties::default(), range: Interval::make_unbounded(&DataType::Null).unwrap(), preserves_lex_ordering: false, + strictly_order_preserving: false, } } @@ -172,4 +227,14 @@ impl ExprProperties { self.preserves_lex_ordering = preserves_lex_ordering; self } + + /// Sets whether the expression is strictly order-preserving and returns + /// the modified instance. + pub fn with_strictly_order_preserving( + mut self, + strictly_order_preserving: bool, + ) -> Self { + self.strictly_order_preserving = strictly_order_preserving; + self + } } diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 4c51ff46f7365..2de3be4c10fa4 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -380,6 +380,11 @@ impl ScalarUDF { self.inner.preserves_lex_ordering(inputs) } + /// See [`ScalarUDFImpl::strictly_order_preserving`] for more details. + pub fn strictly_order_preserving(&self, inputs: &[ExprProperties]) -> Result { + self.inner.strictly_order_preserving(inputs) + } + /// See [`ScalarUDFImpl::coerce_types`] for more details. pub fn coerce_types(&self, arg_types: &[DataType]) -> Result> { self.inner.coerce_types(arg_types) @@ -979,10 +984,20 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { /// Returns true if the function preserves lexicographical ordering based on /// the input ordering. + /// + /// See [`ExprProperties::preserves_lex_ordering`] for more details fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result { Ok(false) } + /// Returns true if the function is strictly order-preserving with respect + /// to its `Ordered` inputs, i.e. `a.cmp(b) == f(a).cmp(f(b))`. + /// + /// See [`ExprProperties::strictly_order_preserving`] for more details + fn strictly_order_preserving(&self, _inputs: &[ExprProperties]) -> Result { + Ok(false) + } + /// Coerce arguments of a function call to types that the function can evaluate. /// /// This function is only called if [`ScalarUDFImpl::signature`] returns @@ -1194,6 +1209,10 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.preserves_lex_ordering(inputs) } + fn strictly_order_preserving(&self, inputs: &[ExprProperties]) -> Result { + self.inner.strictly_order_preserving(inputs) + } + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { self.inner.coerce_types(arg_types) } diff --git a/datafusion/ffi/src/expr/expr_properties.rs b/datafusion/ffi/src/expr/expr_properties.rs index 5b37cc6a28535..584f774c7b26e 100644 --- a/datafusion/ffi/src/expr/expr_properties.rs +++ b/datafusion/ffi/src/expr/expr_properties.rs @@ -29,6 +29,7 @@ pub struct FFI_ExprProperties { sort_properties: FFI_SortProperties, range: FFI_Interval, preserves_lex_ordering: bool, + strictly_order_preserving: bool, } impl TryFrom<&ExprProperties> for FFI_ExprProperties { @@ -41,6 +42,7 @@ impl TryFrom<&ExprProperties> for FFI_ExprProperties { sort_properties, range, preserves_lex_ordering: value.preserves_lex_ordering, + strictly_order_preserving: value.strictly_order_preserving, }) } } @@ -54,6 +56,7 @@ impl TryFrom for ExprProperties { sort_properties, range, preserves_lex_ordering: value.preserves_lex_ordering, + strictly_order_preserving: value.strictly_order_preserving, }) } } diff --git a/datafusion/functions/src/datetime/from_unixtime.rs b/datafusion/functions/src/datetime/from_unixtime.rs index 4787c75b610b6..85494f3abff73 100644 --- a/datafusion/functions/src/datetime/from_unixtime.rs +++ b/datafusion/functions/src/datetime/from_unixtime.rs @@ -22,6 +22,7 @@ use arrow::datatypes::TimeUnit::Second; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err}; use datafusion_expr::TypeSignature::Exact; +use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, @@ -147,6 +148,24 @@ impl ScalarUDFImpl for FromUnixtimeFunc { } } + fn output_ordering(&self, inputs: &[ExprProperties]) -> Result { + // The optional timezone argument must be a constant string and only + // affects the display metadata, not the stored epoch value, so the + // output ordering follows the first argument. + Ok(inputs[0].sort_properties) + } + + fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result { + Ok(true) + } + + fn strictly_order_preserving(&self, _inputs: &[ExprProperties]) -> Result { + // `from_unixtime` stores the input's exact `Int64` value as a + // `Timestamp(Second)`: the mapping is one-to-one, order-preserving, + // and maps nulls to nulls. + Ok(true) + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } diff --git a/datafusion/functions/src/math/monotonicity.rs b/datafusion/functions/src/math/monotonicity.rs index 52449f9c9e0b9..d1174d77b9db1 100644 --- a/datafusion/functions/src/math/monotonicity.rs +++ b/datafusion/functions/src/math/monotonicity.rs @@ -761,6 +761,7 @@ mod tests { .unwrap(), sort_properties: sp, preserves_lex_ordering: false, + strictly_order_preserving: false, } } diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index 2ce8a8d246fe7..15637d24e8a4b 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -329,7 +329,7 @@ mod tests { EquivalenceClass, EquivalenceGroup, EquivalenceProperties, OrderingEquivalenceClass, convert_to_orderings, convert_to_sort_exprs, }; - use crate::expressions::{BinaryExpr, Column, col}; + use crate::expressions::{BinaryExpr, CastExpr, Column, col}; use crate::utils::tests::TestScalarUDF; use crate::{ AcrossPartitions, ConstExpr, PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, @@ -376,6 +376,45 @@ mod tests { Ok(()) } + #[test] + fn test_ordering_satisfy_strictly_order_preserving() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int64, true), + ])); + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let asc = SortOptions::default(); + let sort_a = PhysicalSortExpr::new(Arc::clone(&col_a), asc); + let sort_b = PhysicalSortExpr::new(Arc::clone(&col_b), asc); + let eq_properties = EquivalenceProperties::new_with_orderings( + Arc::clone(&schema), + [vec![sort_a.clone(), sort_b.clone()]], + ); + + assert!(eq_properties.ordering_satisfy(vec![sort_a.clone(), sort_b.clone()])?); + assert!(eq_properties.ordering_satisfy(vec![sort_a.clone()])?); + + // A widening cast is strictly order-preserving: `a` is constant + // within each group of equal `CAST(a AS BIGINT)` values, so `b` + // remains sorted within those groups. + let widening = Arc::new(CastExpr::new(Arc::clone(&col_a), DataType::Int64, None)) + as PhysicalExprRef; + let sort_widening = PhysicalSortExpr::new(widening, asc); + assert!(eq_properties.ordering_satisfy(vec![sort_widening, sort_b.clone()])?); + + // A narrowing cast is only monotonic: it satisfies as a leading key, + // but it may collapse distinct `a` values, so `b` is not guaranteed + // to be sorted within its tie groups. + let narrowing = Arc::new(CastExpr::new(Arc::clone(&col_a), DataType::Int16, None)) + as PhysicalExprRef; + let sort_narrowing = PhysicalSortExpr::new(narrowing, asc); + assert!(eq_properties.ordering_satisfy(vec![sort_narrowing.clone()])?); + assert!(!eq_properties.ordering_satisfy(vec![sort_narrowing, sort_b.clone()])?); + + Ok(()) + } + #[test] fn test_ordering_satisfy_with_equivalence2() -> Result<()> { let test_schema = create_test_schema()?; diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index 17c3898fd9c89..22b3382f50638 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -33,13 +33,13 @@ use self::dependency::{ use crate::equivalence::{ AcrossPartitions, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping, }; -use crate::expressions::{CastExpr, Column, Literal, with_new_schema}; +use crate::expressions::{Column, Literal, with_new_schema}; use crate::{ ConstExpr, LexOrdering, LexRequirement, PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement, }; -use arrow::datatypes::{DataType, SchemaRef}; +use arrow::datatypes::SchemaRef; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Constraint, Constraints, HashMap, Result, plan_err}; use datafusion_expr::interval_arithmetic::Interval; @@ -195,24 +195,30 @@ impl OrderingEquivalenceCache { } impl EquivalenceProperties { - /// Helper used by the ordering equivalence rule when considering whether a - /// cast-bearing expression can replace an existing sort key without - /// invalidating the ordering. + /// Helper used by the ordering equivalence rule when considering whether + /// an expression can replace an existing sort key without invalidating + /// the ordering. /// - /// The substitution is only allowed when the cast wraps the very same child - /// expression that the original sort used and the casted type is a - /// widening/order-preserving conversion. Without those restrictions, a - /// narrowing cast could collapse distinct values and violate the existing + /// The substitution is only allowed when, treating the sort key as the + /// only ordered input, the expression reports the same ordering *and* + /// that it is a one-to-one, order-preserving function of it (see + /// [`ExprProperties::strictly_order_preserving`]). For example, a + /// widening `CAST` of the sort key qualifies, while a narrowing one does + /// not, as it could collapse distinct values and violate the existing /// sort order. - fn substitute_cast_ordering( + fn substitute_order_preserving_ordering( r_expr: Arc, sort_expr: &PhysicalSortExpr, - expr_type: &DataType, + schema: &SchemaRef, ) -> Option { - let cast_expr = r_expr.downcast_ref::()?; - - (cast_expr.expr().eq(&sort_expr.expr) - && CastExpr::check_bigger_cast(cast_expr.cast_type(), expr_type)) + if r_expr.eq(&sort_expr.expr) { + // No point in substituting an expression with itself. + return None; + } + let dependencies = Dependencies::new(std::iter::once(sort_expr.clone())); + let properties = get_expr_properties(&r_expr, &dependencies, schema).ok()?; + (properties.strictly_order_preserving + && properties.sort_properties == SortProperties::Ordered(sort_expr.options)) .then(|| PhysicalSortExpr::new(r_expr, sort_expr.options)) } @@ -482,6 +488,7 @@ impl EquivalenceProperties { sort_properties: SortProperties::Ordered(next.options), range: Interval::make_unbounded(&data_type)?, preserves_lex_ordering: true, + strictly_order_preserving: true, }); } // Check if the expression is monotonic in all arguments: @@ -626,24 +633,55 @@ impl EquivalenceProperties { if !satisfy { return Ok(false); } - // Treat satisfied keys as constants in subsequent iterations. We - // can do this because the "next" key only matters in a lexicographical - // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - // - // For example, assume that the requirement is `[a ASC, (b + c) ASC]`, - // and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. - // From the analysis above, we know that `[a ASC]` is satisfied. Then, - // we add column `a` as constant to the algorithm state. This enables us - // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. - let const_expr = ConstExpr::from(element.expr); - eq_properties.add_constants(std::iter::once(const_expr))?; + // Treat satisfied keys (and the sub-expressions they pin down) as + // constants in subsequent iterations. See + // [`Self::add_satisfied_key_constants`] for the rationale. + eq_properties.add_satisfied_key_constants(element.expr)?; } Ok(true) } + /// Registers a satisfied sort key as a constant for subsequent iterations + /// of the ordering satisfaction checks. We can do this because the "next" + /// key only matters in a lexicographical ordering when the keys to its + /// left have the same values (i.e. within a single tie group). Note that + /// these expressions are not properly "constants"; this is just an + /// implementation strategy confined to the satisfaction checks. + /// + /// For example, assume that the requirement is `[a ASC, (b + c) ASC]`, + /// and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. + /// Once we deduce that `[a ASC]` is satisfied, we add column `a` as a + /// constant to the algorithm state. This enables us to deduce that + /// `(b + c) ASC` is satisfied, given `a` is constant. + /// + /// In addition to the key itself, this also registers any sub-expressions + /// whose values the key pins down: if an expression is strictly + /// order-preserving, equal outputs imply equal values of its ordered + /// children, so within a tie group of the key those children are constant + /// as well. For example, if data is sorted by `[a, b]`, the requirement + /// `[CAST(a AS BIGINT) ASC, b ASC]` is satisfied: `a` is constant within + /// each group of equal `CAST(a AS BIGINT)` values, and hence `b` is + /// sorted within each such group. + fn add_satisfied_key_constants(&mut self, expr: Arc) -> Result<()> { + let mut stack = vec![expr]; + while let Some(expr) = stack.pop() { + let properties = self.get_expr_properties(Arc::clone(&expr)); + if properties.strictly_order_preserving { + for child in expr.children() { + let child_properties = self.get_expr_properties(Arc::clone(child)); + if matches!( + child_properties.sort_properties, + SortProperties::Ordered(_) + ) { + stack.push(Arc::clone(child)); + } + } + } + self.add_constants(std::iter::once(ConstExpr::from(expr)))?; + } + Ok(()) + } + /// Returns the number of consecutive sort expressions (starting from the /// left) that are satisfied by the existing ordering. fn common_sort_prefix_length(&self, normal_ordering: &LexOrdering) -> Result { @@ -676,20 +714,10 @@ impl EquivalenceProperties { // many we've satisfied so far: return Ok(idx); } - // Treat satisfied keys as constants in subsequent iterations. We - // can do this because the "next" key only matters in a lexicographical - // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - // - // For example, assume that the requirement is `[a ASC, (b + c) ASC]`, - // and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. - // From the analysis above, we know that `[a ASC]` is satisfied. Then, - // we add column `a` as constant to the algorithm state. This enables us - // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. - let const_expr = ConstExpr::from(Arc::clone(&element.expr)); - eq_properties.add_constants(std::iter::once(const_expr))? + // Treat satisfied keys (and the sub-expressions they pin down) as + // constants in subsequent iterations. See + // [`Self::add_satisfied_key_constants`] for the rationale. + eq_properties.add_satisfied_key_constants(Arc::clone(&element.expr))?; } // All sort expressions are satisfied, return full length: Ok(full_length) @@ -840,7 +868,9 @@ impl EquivalenceProperties { /// /// TODO: Handle all scenarios that allow substitution; e.g. when `x` is /// sorted, `atan(x + 1000)` should also be substituted. For now, we - /// only consider single-column `CAST` expressions. + /// consider widening `CAST` expressions and single-child expressions + /// that declare themselves one-to-one order-preserving via + /// [`ExprProperties::strictly_order_preserving`]. fn substitute_oeq_class( schema: &SchemaRef, mapping: &ProjectionMapping, @@ -852,21 +882,17 @@ impl EquivalenceProperties { order .into_iter() .map(|sort_expr| { - // The sort expression comes from this schema, so the - // following call to `unwrap` is safe. - let expr_type = sort_expr.expr.data_type(schema).unwrap(); let original_sort_expr = sort_expr.clone(); - // TODO: Add one-to-one analysis for ScalarFunctions. mapping .iter() .map(|(source, _target)| source) .filter(|source| expr_refers(source, &original_sort_expr.expr)) .cloned() .filter_map(|r_expr| { - Self::substitute_cast_ordering( + Self::substitute_order_preserving_ordering( r_expr, &original_sort_expr, - &expr_type, + schema, ) }) .chain(std::iter::once(sort_expr)) @@ -1407,7 +1433,10 @@ fn update_properties( } else if node.expr.is::() { // We have a Column, which is the other possible leaf node type: node.data.range = - Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)? + Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)?; + // A column is the identity mapping of itself, which is trivially + // strict: + node.data.strictly_order_preserving = true; } // Now, check what we know about orderings: let normal_expr = eq_properties @@ -1469,23 +1498,36 @@ fn get_expr_properties( schema: &SchemaRef, ) -> Result { if let Some(column_order) = dependencies.iter().find(|&order| expr.eq(&order.expr)) { - // If exact match is found, return its ordering. + // If exact match is found, return its ordering. This is a base case + // of the recursion: the expression is treated as an atomic ordered + // input from here on, so `strictly_order_preserving` states only that + // it is a one-to-one mapping *of itself* (the identity), which holds + // for any expression. It makes no claim about the expression being + // one-to-one in its own inputs (e.g. `floor(x)` as a sort key), and + // it does not need to: parent expressions are substituted for this + // sort key, so their strictness only has to be relative to it. Ok(ExprProperties { sort_properties: SortProperties::Ordered(column_order.options), range: Interval::make_unbounded(&expr.data_type(schema)?)?, preserves_lex_ordering: false, + strictly_order_preserving: true, }) } else if expr.downcast_ref::().is_some() { Ok(ExprProperties { sort_properties: SortProperties::Unordered, range: Interval::make_unbounded(&expr.data_type(schema)?)?, preserves_lex_ordering: false, + // A base case of the recursion: a column is the identity mapping + // of itself, which is trivially one-to-one. + strictly_order_preserving: true, }) } else if let Some(literal) = expr.downcast_ref::() { Ok(ExprProperties { sort_properties: SortProperties::Singleton, range: literal.value().into(), preserves_lex_ordering: true, + // Vacuously true: a literal has no ordered inputs. + strictly_order_preserving: true, }) } else { // Find orderings of its children diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index 89828620a5930..8b71b3ec409c7 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -764,41 +764,49 @@ impl PhysicalExpr for BinaryExpr { sort_properties: l_order.add(&r_order), range: l_range.add(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::Minus => Ok(ExprProperties { sort_properties: l_order.sub(&r_order), range: l_range.sub(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::Gt => Ok(ExprProperties { sort_properties: l_order.gt_or_gteq(&r_order), range: l_range.gt(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::GtEq => Ok(ExprProperties { sort_properties: l_order.gt_or_gteq(&r_order), range: l_range.gt_eq(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::Lt => Ok(ExprProperties { sort_properties: r_order.gt_or_gteq(&l_order), range: l_range.lt(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::LtEq => Ok(ExprProperties { sort_properties: r_order.gt_or_gteq(&l_order), range: l_range.lt_eq(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::And => Ok(ExprProperties { sort_properties: r_order.and_or(&l_order), range: l_range.and(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::Or => Ok(ExprProperties { sort_properties: r_order.and_or(&l_order), range: l_range.or(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), _ => Ok(ExprProperties::new_unknown()), } diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index 8be2e187d72f7..dbb91e365af90 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -211,8 +211,18 @@ pub(crate) fn cast_expr_properties( target_type: &DataType, ) -> Result { let unbounded = Interval::make_unbounded(target_type)?; - if is_order_preserving_cast_family(&child.range.data_type(), target_type) { - Ok(child.clone().with_range(unbounded)) + let source_type = child.range.data_type(); + // A widening cast is additionally one-to-one, so it is strictly + // order-preserving; a narrowing cast may collapse distinct values, + // breaking the ordering of subsequent sort keys. + let bigger_cast = CastExpr::check_bigger_cast(target_type, &source_type); + if is_order_preserving_cast_family(&source_type, target_type) || bigger_cast { + Ok(child + .clone() + .with_range(unbounded) + .with_strictly_order_preserving( + child.strictly_order_preserving && bigger_cast, + )) } else { Ok(ExprProperties::new_unknown().with_range(unbounded)) } diff --git a/datafusion/physical-expr/src/expressions/literal.rs b/datafusion/physical-expr/src/expressions/literal.rs index 5fb9a3b2cd29b..a7af824230780 100644 --- a/datafusion/physical-expr/src/expressions/literal.rs +++ b/datafusion/physical-expr/src/expressions/literal.rs @@ -123,6 +123,8 @@ impl PhysicalExpr for Literal { sort_properties: SortProperties::Singleton, range: Interval::try_new(self.value().clone(), self.value().clone())?, preserves_lex_ordering: true, + // Vacuously true: a literal has no ordered inputs. + strictly_order_preserving: true, }) } diff --git a/datafusion/physical-expr/src/expressions/negative.rs b/datafusion/physical-expr/src/expressions/negative.rs index 9fbf38361c89c..c894c12784dc5 100644 --- a/datafusion/physical-expr/src/expressions/negative.rs +++ b/datafusion/physical-expr/src/expressions/negative.rs @@ -166,6 +166,8 @@ impl PhysicalExpr for NegativeExpr { sort_properties: -children[0].sort_properties, range: children[0].range.clone().arithmetic_negate()?, preserves_lex_ordering: false, + // Negation is one-to-one but reverses the ordering direction. + strictly_order_preserving: false, }) } diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 418d005c971ea..6a5ab219aa8dd 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -316,6 +316,7 @@ impl PhysicalExpr for ScalarFunctionExpr { fn get_properties(&self, children: &[ExprProperties]) -> Result { let sort_properties = self.fun.output_ordering(children)?; let preserves_lex_ordering = self.fun.preserves_lex_ordering(children)?; + let strictly_order_preserving = self.fun.strictly_order_preserving(children)?; let children_range = children .iter() .map(|props| &props.range) @@ -326,6 +327,7 @@ impl PhysicalExpr for ScalarFunctionExpr { sort_properties, range, preserves_lex_ordering, + strictly_order_preserving, }) } diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index 978fcc197c0de..a267ddddddd54 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -709,6 +709,104 @@ physical_plan statement ok drop table multiple_ordered_table; + +# Create a table having dependent sort order +statement ok +CREATE EXTERNAL TABLE multiple_ordered_table ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER +) +STORED AS CSV +WITH ORDER (a ASC, b ASC, c ASC) +LOCATION '../core/tests/data/window_2.csv' +OPTIONS ('format.has_header' 'true'); + +# Test without repartition so removal of sort is more apperant +statement ok +set datafusion.execution.target_partitions = 1; + +# A strictly order-preserving scalar function is one-to-one, so an ordering on +# its argument carries over to its result. `from_unixtime` reinterprets the +# input integer as a timestamp without changing the value, so the whole +# ordering is preserved and no SortExec is needed. +query TT +EXPLAIN SELECT from_unixtime(a) AS a_, from_unixtime(b) AS b_, from_unixtime(c) AS c_ +FROM multiple_ordered_table +ORDER BY a_, b_, c_; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, b_ ASC NULLS LAST, c_ ASC NULLS LAST +02)--Projection: from_unixtime(CAST(multiple_ordered_table.a AS Int64)) AS a_, from_unixtime(CAST(multiple_ordered_table.b AS Int64)) AS b_, from_unixtime(CAST(multiple_ordered_table.c AS Int64)) AS c_ +03)----TableScan: multiple_ordered_table projection=[a, b, c] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[from_unixtime(CAST(a@1 AS Int64)) as a_, from_unixtime(CAST(b@2 AS Int64)) as b_, from_unixtime(CAST(c@3 AS Int64)) as c_], file_type=csv, has_header=true + +# Being one-to-one also justifies keeping the *suffix* sort keys: data sorted +# by [a, b] is also sorted by [from_unixtime(a), b], because rows with equal +# `a_` have equal `a`, within which `b` is already sorted. +query TT +EXPLAIN SELECT from_unixtime(a) AS a_, b +FROM multiple_ordered_table +ORDER BY a_, b; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +02)--Projection: from_unixtime(CAST(multiple_ordered_table.a AS Int64)) AS a_, multiple_ordered_table.b +03)----TableScan: multiple_ordered_table projection=[a, b] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[from_unixtime(CAST(a@1 AS Int64)) as a_, b], file_type=csv, has_header=true + +# A widening CAST is one-to-one too: +query TT +EXPLAIN SELECT CAST(a AS BIGINT) AS a_, b +FROM multiple_ordered_table +ORDER BY a_, b; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +02)--Projection: CAST(multiple_ordered_table.a AS Int64) AS a_, multiple_ordered_table.b +03)----TableScan: multiple_ordered_table projection=[a, b] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[CAST(a@1 AS Int64) as a_, b], file_type=csv, has_header=true + +# In contrast, a merely monotone (`preserves_lex_ordering`, but not strictly +# order-preserving) function such as floor() does NOT justify the suffix keys: +# in general floor() collapses distinct inputs into one output value, and `b` +# is not sorted within such a run, so a SortExec must remain. +query TT +EXPLAIN SELECT floor(CAST(a AS DOUBLE)) AS a_, b +FROM multiple_ordered_table +ORDER BY a_, b; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +02)--Projection: floor(CAST(multiple_ordered_table.a AS Float64)) AS a_, multiple_ordered_table.b +03)----TableScan: multiple_ordered_table projection=[a, b] +physical_plan +01)SortExec: expr=[a_@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[floor(CAST(a@1 AS Float64)) as a_, b], file_type=csv, has_header=true + +# Monotonicity alone is still enough when the expression is the *only* sort +# key, so here the SortExec is removed even though floor() is not strict: +query TT +EXPLAIN SELECT floor(CAST(a AS DOUBLE)) AS a_ +FROM multiple_ordered_table +ORDER BY a_; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST +02)--Projection: floor(CAST(multiple_ordered_table.a AS Float64)) AS a_ +03)----TableScan: multiple_ordered_table projection=[a] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[floor(CAST(a@1 AS Float64)) as a_], file_type=csv, has_header=true + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# reset it explicitly. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +drop table multiple_ordered_table; + # Create tables having some ordered columns. In the next step, we will expect to observe that scalar # functions, such as mathematical functions like atan(), ceil(), sqrt(), or date_time functions # like date_bin() and date_trunc(), will maintain the order of its argument columns.