ScalarUdfImpl::strictly_order_preserving: Allow expression to report whether they keep the same ordering of the input

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<dyn PhysicalExpr>,
+}
+
+impl AssertNotNull {
+    fn new(inner: Arc<dyn PhysicalExpr>) -> Arc<Self> {
+        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<H: std::hash::Hasher>(&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<DataType> {
+        self.inner.data_type(input_schema)
+    }
+
+    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+        Ok(false)
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        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<dyn PhysicalExpr>> {
+        vec![&self.inner]
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        Ok(Arc::new(AssertNotNull {
+            inner: Arc::clone(&children[0]),
+        }))
+    }
+
+    fn get_properties(
+        &self,
+        children: &[datafusion_expr::sort_properties::ExprProperties],
+    ) -> Result<datafusion_expr::sort_properties::ExprProperties> {
+        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<dyn PhysicalExpr>, String)>,
+        input: Arc<dyn ExecutionPlan>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let proj_exprs: Vec<ProjectionExpr> = 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<dyn ExecutionPlan>;
+
+    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(())
+}

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
+    }
 }

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<bool> {
+        self.inner.strictly_order_preserving(inputs)
+    }
+
     /// See [`ScalarUDFImpl::coerce_types`] for more details.
     pub fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
         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<bool> {
         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<bool> {
+        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<bool> {
+        self.inner.strictly_order_preserving(inputs)
+    }
+
     fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> {
         self.inner.coerce_types(arg_types)
     }

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<FFI_ExprProperties> for ExprProperties {
             sort_properties,
             range,
             preserves_lex_ordering: value.preserves_lex_ordering,
+            strictly_order_preserving: value.strictly_order_preserving,
         })
     }
 }

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<SortProperties> {
+        // 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<bool> {
+        Ok(true)
+    }
+
+    fn strictly_order_preserving(&self, _inputs: &[ExprProperties]) -> Result<bool> {
+        // `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()
     }

datafusion/functions/src/math/monotonicity.rs

@@ -761,6 +761,7 @@ mod tests {
                 .unwrap(),
                 sort_properties: sp,
                 preserves_lex_ordering: false,
+                strictly_order_preserving: false,
             }
         }