Address code review comments

julian-elastic · julian-elastic · commit c4f70d690a67 · 2025-07-29T21:32:56.000-04:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/FunctionUtils.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/FunctionUtils.java
@@ -8,6 +8,7 @@
 
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Strings;
 import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.common.Failure;
 import org.elasticsearch.xpack.esql.common.Failures;
@@ -148,4 +149,13 @@ public static String queryAsString(Expression queryField, String sourceText) {
             format(null, "Query value must be a constant string in [{}], found [{}]", sourceText, queryField)
         );
     }
+
+    public static int intValueOf(Expression field, String sourceText, String fieldName) {
+        if (field instanceof Literal literal) {
+            return ((Number) literal.value()).intValue();
+        }
+        throw new EsqlIllegalArgumentException(
+            Strings.format(null, "[{}] value must be a constant number in [{}], found [{}]", fieldName, sourceText, field)
+        );
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/CountDistinct.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/CountDistinct.java
@@ -42,13 +42,13 @@
 import java.util.Map;
 import java.util.function.Function;
 
-import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isWholeNumber;
 import static org.elasticsearch.xpack.esql.core.util.CollectionUtils.nullSafeList;
+import static org.elasticsearch.xpack.esql.expression.function.FunctionUtils.intValueOf;
 
 public class CountDistinct extends AggregateFunction implements OptionalArgument, ToAggregator, SurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
@@ -219,12 +219,7 @@ public AggregatorFunctionSupplier supplier() {
     }
 
     private int precisionValue() {
-        if (precision instanceof Literal literal) {
-            return ((Number) literal.value()).intValue();
-        }
-        throw new EsqlIllegalArgumentException(
-            format(null, "Precision value must be a constant integer in [{}], found [{}]", source(), precision)
-        );
+        return intValueOf(precision, source().text(), "Precision");
     }
 
     @Override
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Percentile.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Percentile.java
@@ -15,7 +15,6 @@
 import org.elasticsearch.compute.aggregation.PercentileDoubleAggregatorFunctionSupplier;
 import org.elasticsearch.compute.aggregation.PercentileIntAggregatorFunctionSupplier;
 import org.elasticsearch.compute.aggregation.PercentileLongAggregatorFunctionSupplier;
-import org.elasticsearch.xpack.esql.EsqlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
@@ -34,11 +33,11 @@
 import java.util.List;
 
 import static java.util.Collections.singletonList;
-import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
+import static org.elasticsearch.xpack.esql.expression.function.FunctionUtils.intValueOf;
 
 public class Percentile extends NumericAggregate implements SurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
@@ -170,12 +169,7 @@ protected AggregatorFunctionSupplier doubleSupplier() {
     }
 
     private int percentileValue() {
-        if (percentile() instanceof Literal literal) {
-            return ((Number) literal.value()).intValue();
-        }
-        throw new EsqlIllegalArgumentException(
-            format(null, "Percentile value must be a constant number in [{}], found [{}]", source(), percentile)
-        );
+        return intValueOf(percentile(), source().text(), "Percentile");
     }
 
     @Override
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/CIDRMatch.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ip/CIDRMatch.java
@@ -189,7 +189,7 @@ public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHand
         Check.isTrue(Expressions.foldable(matches), "Expected foldable matches, but got [{}]", matches);
 
         String targetFieldName = handler.nameOf(fa.exactAttribute());
-        Set<Object> set = new LinkedHashSet<>(matches.stream().map(Foldables::extractLiteralOrReturnSelf).toList());
+        Set<Object> set = new LinkedHashSet<>(matches.stream().map(Foldables::literalValueOf).toList());
 
         return new TermsQuery(source(), targetFieldName, set);
     }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSort.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/multivalue/MvSort.java
@@ -51,7 +51,6 @@
 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;
@@ -308,18 +307,6 @@ 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;
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java
@@ -615,16 +615,17 @@ public void testReplaceStatsFilteredAggWithEvalFilterFalseAndNull() {
             from test
             | stats sum(salary) + 1 where false,
                     sum(salary) + 3,
-                    sum(salary) + 2 where null
+                    sum(salary) + 2 where null,
+                    sum(salary) + 4 where not true
             """);
 
         var project = as(plan, Project.class);
         assertThat(
             Expressions.names(project.projections()),
-            contains("sum(salary) + 1 where false", "sum(salary) + 3", "sum(salary) + 2 where null")
+            contains("sum(salary) + 1 where false", "sum(salary) + 3", "sum(salary) + 2 where null", "sum(salary) + 4 where not true")
         );
         var eval = as(project.child(), Eval.class);
-        assertThat(eval.fields().size(), is(3));
+        assertThat(eval.fields().size(), is(4));
 
         var alias = as(eval.fields().getFirst(), Alias.class);
         assertTrue(alias.child().foldable());
@@ -664,6 +665,42 @@ public void testReplaceStatsFilteredAggWithEvalCount() {
         assertThat(block.asVector().getLong(0), is(0L));
     }
 
+    /*
+     * Limit[1000[INTEGER]]
+     * \_LocalRelation[[count(salary) where false{r}#3],[LongVectorBlock[vector=ConstantLongVector[positions=1, value=0]]]]
+     */
+    public void testReplaceStatsFilteredAggWithEvalNotTrue() {
+        var plan = plan("""
+            from test
+            | stats count(salary) where not true
+            """);
+
+        var limit = as(plan, Limit.class);
+        var source = as(limit.child(), LocalRelation.class);
+        assertThat(Expressions.names(source.output()), contains("count(salary) where not true"));
+        Block[] blocks = source.supplier().get();
+        assertThat(blocks.length, is(1));
+        var block = as(blocks[0], LongVectorBlock.class);
+        assertThat(block.getPositionCount(), is(1));
+        assertThat(block.asVector().getLong(0), is(0L));
+    }
+
+    /*
+    * Limit[1000[INTEGER],false]
+    * \_Aggregate[[],[COUNT(salary{f}#10,true[BOOLEAN]) AS m1#4]]
+    * \_EsRelation[test][_meta_field{f}#11, emp_no{f}#5, first_name{f}#6, ge..]
+    */
+    public void testReplaceStatsFilteredAggWithEvalNotFalse() {
+        var plan = plan("""
+            from test
+            | stats m1 = count(salary) where not false
+            """);
+        var limit = as(plan, Limit.class);
+        var aggregate = as(limit.child(), Aggregate.class);
+        assertThat(Expressions.names(aggregate.aggregates()), contains("m1"));
+        var source = as(aggregate.child(), EsRelation.class);
+    }
+
     /*
      * Project[[count_distinct(salary + 2) + 3 where false{r}#3]]
      * \_Eval[[$$COUNTDISTINCT$count_distinct(>$0{r$}#15 + 3[INTEGER] AS count_distinct(salary + 2) + 3 where false]]
@@ -815,6 +852,37 @@ public void testExtractStatsCommonFilter() {
         var source = as(filter.child(), EsRelation.class);
     }
 
+    /*
+     Limit[1000[INTEGER],false]
+     \_Aggregate[[],[MIN(salary{f}#15,true[BOOLEAN]) AS m1#4, MAX(salary{f}#15,true[BOOLEAN]) AS m2#8]]
+     \_Filter[emp_no{f}#10 > 1[INTEGER]]
+     \_EsRelation[test][_meta_field{f}#16, emp_no{f}#10, first_name{f}#11, ..]
+     */
+    public void testExtractStatsCommonAlwaysTruePlusOtherFilter() {
+        var plan = plan("""
+            from test
+            | stats m1 = min(salary) where (true and emp_no > 1),
+                    m2 = max(salary) where (1==1 and emp_no > 1)
+            """);
+
+        var limit = as(plan, Limit.class);
+        var agg = as(limit.child(), Aggregate.class);
+        assertThat(agg.aggregates().size(), is(2));
+
+        var alias = as(agg.aggregates().get(0), Alias.class);
+        var aggFunc = as(alias.child(), AggregateFunction.class);
+        assertThat(aggFunc.filter(), is(Literal.TRUE));
+
+        alias = as(agg.aggregates().get(1), Alias.class);
+        aggFunc = as(alias.child(), AggregateFunction.class);
+        assertThat(aggFunc.filter(), is(Literal.TRUE));
+
+        var filter = as(agg.child(), Filter.class);
+        assertThat(Expressions.name(filter.condition()), is("emp_no > 1"));
+
+        var source = as(filter.child(), EsRelation.class);
+    }
+
     public void testExtractStatsCommonFilterUsingAliases() {
         var plan = plan("""
             from test
diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/230_folding.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/230_folding.yml
@@ -26,18 +26,20 @@ setup:
                 similarity: l2_norm
               ip:
                 type: ip
+              favorite_numbers:
+                type: integer
 
   - do:
       bulk:
         index: employees
         refresh: true
         body:
           - { "index": { } }
-          - { "hire_date": "2020-01-01", "salary_change": 100.5, "salary": 50000, "salary_change_long": 100, "name": "Alice Smith", "image_vector": [ 0.1, 0.2, 0.3 ], "ip": "192.168.1.1" }
+          - { "hire_date": "2020-01-01", "salary_change": 100.5, "salary": 50000, "salary_change_long": 100, "name": "Alice Smith", "image_vector": [ 0.1, 0.2, 0.3 ], "ip": "192.168.1.1" , "favorite_numbers": [ 4, 2, 3 ] }
           - { "index": { } }
-          - { "hire_date": "2021-01-01", "salary_change": 200.5, "salary": 60000, "salary_change_long": 200, "name": "Bob Johnson", "image_vector": [ 0.4, 0.5, 0.6 ], "ip": "192.168.1.2" }
+          - { "hire_date": "2021-01-01", "salary_change": 200.5, "salary": 60000, "salary_change_long": 200, "name": "Bob Johnson", "image_vector": [ 0.4, 0.5, 0.6 ], "ip": "192.168.1.2" , "favorite_numbers": [ 4, 7, -2 ] }
           - { "index": { } }
-          - { "hire_date": "2019-01-01", "salary_change": 50.5, "salary": 40000, "salary_change_long": 50, "name": "Charlie Smith", "image_vector": [ 0.7, 0.8, 0.9 ], "ip": "10.168.1.2" }
+          - { "hire_date": "2019-01-01", "salary_change": 50.5, "salary": 40000, "salary_change_long": 50, "name": "Charlie Smith", "image_vector": [ 0.7, 0.8, 0.9 ], "ip": "10.168.1.2" , "favorite_numbers": [ 2, 1, 3 ] }
 
 ---
 TOP function with constant folding:
@@ -631,6 +633,20 @@ MV_SORT with foldable order:
   - match: { columns.2.name: "sd" }
   - match: { values.0.1: [ -3, 2, 2, 4 ] }
   - match: { values.0.2: [ 4, 2, 2, -3 ] }
+---
+MV_SORT with foldable order on employee favorite numbers:
+  - do:
+      esql.query:
+        body:
+          query: |
+            FROM employees
+            | WHERE salary == 50000
+            | EVAL sa = mv_sort(favorite_numbers, CONCAT("AS", "C"))
+            | KEEP sa
+            | LIMIT 1
+
+  - match: { columns.0.name: "sa" }
+  - match: { values.0.0: [ 2,3,4 ] }
 
 ---
 MV_SORT with invalid foldable order:

Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,7 @@ public Query asQuery(LucenePushdownPredicates pushdownPredicates, TranslatorHand`
`189`	`189`	`Check.isTrue(Expressions.foldable(matches), "Expected foldable matches, but got [{}]", matches);`
`190`	`190`
`191`	`191`	`String targetFieldName = handler.nameOf(fa.exactAttribute());`
`192`		`- Set<Object> set = new LinkedHashSet<>(matches.stream().map(Foldables::extractLiteralOrReturnSelf).toList());`
	`192`	`+ Set<Object> set = new LinkedHashSet<>(matches.stream().map(Foldables::literalValueOf).toList());`
`193`	`193`
`194`	`194`	`return new TermsQuery(source(), targetFieldName, set);`
`195`	`195`	`}`