From d7d8270e39508dc6bc84b8f4964727e544cd5209 Mon Sep 17 00:00:00 2001
From: Goksel Kabadayi <45314116+gokselk@users.noreply.github.com>
Date: Thu, 7 Nov 2024 11:02:14 +0300
Subject: [PATCH] Support n-ary monotonic functions in `discover_new_orderings`

---
 .../src/equivalence/properties.rs             | 64 +++++++++++--------
 1 file changed, 37 insertions(+), 27 deletions(-)
 mode change 100644 => 100755 datafusion/physical-expr/src/equivalence/properties.rs

diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs
old mode 100644
new mode 100755
index 9a16b205ae25..6ae9fde0d03d
--- a/datafusion/physical-expr/src/equivalence/properties.rs
+++ b/datafusion/physical-expr/src/equivalence/properties.rs
@@ -346,40 +346,50 @@ impl EquivalenceProperties {
             .unwrap_or_else(|| vec![Arc::clone(&normalized_expr)]);
 
         let mut new_orderings: Vec<LexOrdering> = vec![];
-        for (ordering, next_expr) in self
+        for ordering in self
             .normalized_oeq_class()
             .iter()
             .filter(|ordering| ordering[0].expr.eq(&normalized_expr))
-            // First expression after leading ordering
-            .filter_map(|ordering| Some(ordering).zip(ordering.get(1)))
         {
             let leading_ordering = ordering[0].options;
-            // Currently, we only handle expressions with a single child.
-            // TODO: It should be possible to handle expressions orderings like
-            //       f(a, b, c), a, b, c if f is monotonic in all arguments.
+
+            // Handle expressions with multiple children
             for equivalent_expr in &eq_class {
                 let children = equivalent_expr.children();
-                if children.len() == 1
-                    && children[0].eq(&next_expr.expr)
-                    && SortProperties::Ordered(leading_ordering)
-                        == equivalent_expr
-                            .get_properties(&[ExprProperties {
-                                sort_properties: SortProperties::Ordered(
-                                    leading_ordering,
-                                ),
-                                range: Interval::make_unbounded(
-                                    &equivalent_expr.data_type(&self.schema)?,
-                                )?,
-                            }])?
-                            .sort_properties
-                {
-                    // Assume existing ordering is [a ASC, b ASC]
-                    // When equality a = f(b) is given, If we know that given ordering `[b ASC]`, ordering `[f(b) ASC]` is valid,
-                    // then we can deduce that ordering `[b ASC]` is also valid.
-                    // Hence, ordering `[b ASC]` can be added to the state as valid ordering.
-                    // (e.g. existing ordering where leading ordering is removed)
-                    new_orderings.push(ordering[1..].to_vec());
-                    break;
+                if children.is_empty() {
+                    continue;
+                }
+
+                // Check if all children match the next expressions in the ordering
+                let mut all_children_match = true;
+                let mut child_properties = vec![];
+
+                // Build properties for each child based on the next expressions
+                for (i, child) in children.iter().enumerate() {
+                    if let Some(next) = ordering.get(i + 1) {
+                        if !child.eq(&next.expr) {
+                            all_children_match = false;
+                            break;
+                        }
+                        child_properties.push(ExprProperties {
+                            sort_properties: SortProperties::Ordered(next.options),
+                            range: Interval::make_unbounded(&child.data_type(&self.schema)?)?,
+                        });
+                    } else {
+                        all_children_match = false;
+                        break;
+                    }
+                }
+
+                if all_children_match {
+                    // Check if the expression is monotonic in all arguments
+                    if let Ok(expr_properties) = equivalent_expr.get_properties(&child_properties) {
+                        if SortProperties::Ordered(leading_ordering) == expr_properties.sort_properties {
+                            // Add ordering with leading expression removed
+                            new_orderings.push(ordering[1..].to_vec());
+                            break;
+                        }
+                    }
                 }
             }
         }