Remove folding from MvSort

Author: julian-elastic

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSort.java

@@ -48,8 +48,10 @@
 import org.elasticsearch.xpack.esql.planner.PlannerUtils;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
 
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
@@ -146,6 +148,19 @@ public boolean foldable() {
         return field.foldable() && (order == null || order.foldable());
     }
 
+    @Override
+    public Object fold(FoldContext ctx) {
+        if (order instanceof Literal == false) {
+            // if we are here, we know that order is foldable, so we can safely fold it
+            Literal newOrder = Literal.of(ctx, order);
+            List<Expression> newChildren = new ArrayList<>();
+            newChildren.add(field);
+            newChildren.add(newOrder);
+            return replaceChildren(newChildren).fold(ctx);
+        }
+        return super.fold(ctx);
+    }
+
     @Override
     public EvalOperator.ExpressionEvaluator.Factory toEvaluator(ToEvaluator toEvaluator) {
         boolean ordering = true;
@@ -246,7 +261,7 @@ public void postOptimizationVerification(Failures failures) {
                     sourceText(),
                     BytesRefs.toString(ASC.value()),
                     BytesRefs.toString(DESC.value()),
-                    BytesRefs.toString(order.fold(FoldContext.small() /* TODO remove me */))
+                    BytesRefs.toString(order)
                 )
             );
         }
@@ -255,16 +270,21 @@ public void postOptimizationVerification(Failures failures) {
     private boolean isValidOrder() {
         boolean isValidOrder = true;
         if (order != null && order.foldable()) {
-            Object obj = order.fold(FoldContext.small() /* TODO remove me */);
-            String o = null;
-            if (obj instanceof BytesRef ob) {
-                o = ob.utf8ToString();
-            } else if (obj instanceof String os) {
-                o = os;
-            }
-            if (o == null
-                || o.equalsIgnoreCase(BytesRefs.toString(ASC.value())) == false
-                    && o.equalsIgnoreCase(BytesRefs.toString(DESC.value())) == false) {
+            if (order instanceof Literal literal) {
+                Object obj = literal.value();
+                String o = null;
+                if (obj instanceof BytesRef ob) {
+                    o = ob.utf8ToString();
+                } else if (obj instanceof String os) {
+                    o = os;
+                }
+                if (o == null
+                    || o.equalsIgnoreCase(BytesRefs.toString(ASC.value())) == false
+                        && o.equalsIgnoreCase(BytesRefs.toString(DESC.value())) == false) {
+                    isValidOrder = false;
+                }
+            } else {
+                // order should be folded already, so if it is not literal it is invalid
                 isValidOrder = false;
             }
         }
@@ -288,6 +308,18 @@ public String toString() {
         }
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (o == null || getClass() != o.getClass()) return false;
+        MvSort mvSort = (MvSort) o;
+        return Objects.equals(field(), mvSort.field()) && Objects.equals(order(), mvSort.order());
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(field(), order());
+    }
+
     private static class Evaluator implements EvalOperator.ExpressionEvaluator {
         private final BlockFactory blockFactory;
         private final EvalOperator.ExpressionEvaluator field;

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/230_folding.yml

@@ -615,3 +615,42 @@ COUNT_DISTINCT with non-foldable precision:
             | STATS COUNT_DISTINCT(salary, salary + 1)  | LIMIT 1
   - match: { error.type: "verification_exception" }
   - contains: { error.reason: "second argument of [COUNT_DISTINCT(salary, salary + 1)] must be a constant, received [salary + 1]" }
+
+---
+MV_SORT with foldable order:
+  - do:
+      esql.query:
+        body:
+          query: |
+            ROW a = [4, 2, -3, 2]
+            | EVAL sa = mv_sort(a, CONCAT("AS", "C")), sd = mv_sort(a, CONCAT("DE", "SC")) | LIMIT 1
+  - match: { columns.0.name: "a" }
+  - match: { columns.1.name: "sa" }
+  - match: { columns.2.name: "sd" }
+  - match: { values.0.1: [ -3, 2, 2, 4 ] }
+  - match: { values.0.2: [ 4, 2, 2, -3 ] }
+
+---
+MV_SORT with invalid foldable order:
+  - do:
+      catch: bad_request
+      esql.query:
+        body:
+          query: |
+            ROW a = [4, 2, -3, 2]
+            | EVAL sa = mv_sort(a, CONCAT("ASC", "DESC")) | LIMIT 1
+  - match: { error.type: "illegal_argument_exception" }
+  - contains: { error.reason: "Invalid order value in [mv_sort(a, CONCAT(\"ASC\", \"DESC\"))], expected one of [ASC, DESC] but got [ASCDESC]" }
+
+
+---
+MV_SORT with invalid non-foldable order:
+  - do:
+      catch: bad_request
+      esql.query:
+        body:
+          query: |
+            ROW a = [4, 2, -3, 2], str = "REVERSE"
+            | EVAL sa = mv_sort(a, CONCAT(str, "DESC")) | LIMIT 1
+  - match: { error.type: "illegal_argument_exception" }
+  - contains: { error.reason: "Invalid order value in [mv_sort(a, CONCAT(str, \"DESC\"))], expected one of [ASC, DESC] but got [REVERSEDESC]" }