Merge branch 'main' into esql-sample-parameter

Author: jan-elastic

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -195,6 +195,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ML_INFERENCE_MISTRAL_CHAT_COMPLETION_ADDED_8_19 = def(8_841_0_47);
     public static final TransportVersion ML_INFERENCE_ELASTIC_RERANK_ADDED_8_19 = def(8_841_0_48);
     public static final TransportVersion NONE_CHUNKING_STRATEGY_8_19 = def(8_841_0_49);
+    public static final TransportVersion IDP_CUSTOM_SAML_ATTRIBUTES_ALLOW_LIST_8_19 = def(8_841_0_50);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);

server/src/main/java/org/elasticsearch/index/shard/IndexShard.java

@@ -1326,12 +1326,15 @@ private Engine.GetResult innerGet(Engine.Get get, boolean translogOnly, Function
     public Engine.RefreshResult refresh(String source) {
         verifyNotClosed();
         logger.trace("refresh with source [{}]", source);
-        return getEngine().refresh(source);
+        return withEngine(engine -> engine.refresh(source));
     }
 
     public void externalRefresh(String source, ActionListener<Engine.RefreshResult> listener) {
         verifyNotClosed();
-        getEngine().externalRefresh(source, listener);
+        withEngine(engine -> {
+            engine.externalRefresh(source, listener);
+            return null;
+        });
     }
 
     /**

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/expression/MetadataAttribute.java

@@ -69,53 +69,13 @@ public MetadataAttribute(Source source, String name, DataType dataType, boolean
         this(source, name, dataType, Nullability.TRUE, null, false, searchable);
     }
 
-    @Deprecated
-    /**
-     * Old constructor from when this had a qualifier string. Still needed to not break serialization.
-     */
-    private MetadataAttribute(
-        Source source,
-        String name,
-        DataType dataType,
-        @Nullable String qualifier,
-        Nullability nullability,
-        @Nullable NameId id,
-        boolean synthetic,
-        boolean searchable
-    ) {
-        this(source, name, dataType, nullability, id, synthetic, searchable);
-    }
-
-    @SuppressWarnings("unchecked")
-    private MetadataAttribute(StreamInput in) throws IOException {
-        /*
-         * The funny casting dance with `(StreamInput & PlanStreamInput) in` is required
-         * because we're in esql-core here and the real PlanStreamInput is in
-         * esql-proper. And because NamedWriteableRegistry.Entry needs StreamInput,
-         * not a PlanStreamInput. And we need PlanStreamInput to handle Source
-         * and NameId. This should become a hard cast when we move everything out
-         * of esql-core.
-         */
-        this(
-            Source.readFrom((StreamInput & PlanStreamInput) in),
-            in.readString(),
-            DataType.readFrom(in),
-            in.readOptionalString(),
-            in.readEnum(Nullability.class),
-            NameId.readFrom((StreamInput & PlanStreamInput) in),
-            in.readBoolean(),
-            in.readBoolean()
-        );
-    }
-
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         if (((PlanStreamOutput) out).writeAttributeCacheHeader(this)) {
             Source.EMPTY.writeTo(out);
             out.writeString(name());
             dataType().writeTo(out);
-            // We used to write the qualifier here. We can still do if needed in the future.
-            out.writeOptionalString(null);
+            out.writeOptionalString(null); // qualifier, no longer used
             out.writeEnum(nullable());
             id().writeTo(out);
             out.writeBoolean(synthetic());
@@ -124,7 +84,25 @@ public void writeTo(StreamOutput out) throws IOException {
     }
 
     public static MetadataAttribute readFrom(StreamInput in) throws IOException {
-        return ((PlanStreamInput) in).readAttributeWithCache(MetadataAttribute::new);
+        /*
+         * The funny casting dance with `(StreamInput & PlanStreamInput) in` is required
+         * because we're in esql-core here and the real PlanStreamInput is in
+         * esql-proper. And because NamedWriteableRegistry.Entry needs StreamInput,
+         * not a PlanStreamInput. And we need PlanStreamInput to handle Source
+         * and NameId. This should become a hard cast when we move everything out
+         * of esql-core.
+         */
+        return ((PlanStreamInput) in).readAttributeWithCache(stream -> {
+            Source source = Source.readFrom((StreamInput & PlanStreamInput) stream);
+            String name = stream.readString();
+            DataType dataType = DataType.readFrom(stream);
+            String qualifier = stream.readOptionalString(); // qualifier, no longer used
+            Nullability nullability = stream.readEnum(Nullability.class);
+            NameId id = NameId.readFrom((StreamInput & PlanStreamInput) stream);
+            boolean synthetic = stream.readBoolean();
+            boolean searchable = stream.readBoolean();
+            return new MetadataAttribute(source, name, dataType, nullability, id, synthetic, searchable);
+        });
     }
 
     @Override
@@ -134,7 +112,7 @@ public String getWriteableName() {
 
     @Override
     protected MetadataAttribute clone(Source source, String name, DataType type, Nullability nullability, NameId id, boolean synthetic) {
-        return new MetadataAttribute(source, name, type, null, nullability, id, synthetic, searchable);
+        return new MetadataAttribute(source, name, type, nullability, id, synthetic, searchable);
     }
 
     @Override

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestSampleTestCase.java

@@ -33,7 +33,7 @@ public class RestSampleTestCase extends ESRestTestCase {
     public void skipWhenSampleDisabled() throws IOException {
         assumeTrue(
             "Requires SAMPLE capability",
-            EsqlSpecTestCase.hasCapabilities(adminClient(), List.of(EsqlCapabilities.Cap.SAMPLE_V2.capabilityName()))
+            EsqlSpecTestCase.hasCapabilities(adminClient(), List.of(EsqlCapabilities.Cap.SAMPLE_V3.capabilityName()))
         );
     }
 

x-pack/plugin/esql/qa/testFixtures/src/main/resources/sample.csv-spec

@@ -9,7 +9,7 @@
 // because the CSV tests don't support such assertions.
 
 row
-required_capability: sample_v2
+required_capability: sample_v3
 
 ROW x = 1 | SAMPLE .999999999
 ;
@@ -20,7 +20,7 @@ x:integer
 
 
 row and mv_expand
-required_capability: sample_v2
+required_capability: sample_v3
 
 ROW x = [1,2,3,4,5] | MV_EXPAND x | SAMPLE .999999999
 ;
@@ -35,15 +35,14 @@ x:integer
 
 
 adjust stats for sampling
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SAMPLE 0.5
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no)), avg_emp_no = AVG(emp_no), sum_emp_no = SUM(emp_no)
-  | EVAL is_expected = count >= 20 AND count <= 180 AND
-                       values_count >= 10 AND values_count <= 90 AND
+  | STATS count = COUNT(), avg_emp_no = AVG(emp_no), sum_emp_no = SUM(emp_no)
+  | EVAL is_expected = count >= 10 AND count <= 90 AND
                        avg_emp_no > 10010 AND avg_emp_no < 10090 AND
-                       sum_emp_no > 20*10010 AND sum_emp_no < 180*10090
+                       sum_emp_no > 10*10010 AND sum_emp_no < 90*10090
   | KEEP is_expected
 ;
 
@@ -53,14 +52,13 @@ true
 
 
 before where
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SAMPLE 0.5
   | WHERE emp_no > 10050 
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no)), avg_emp_no = AVG(emp_no)
-  | EVAL is_expected = count >= 5 AND count <= 95 AND
-                       values_count >= 2 AND values_count <= 48 AND
+  | STATS count = COUNT(), avg_emp_no = AVG(emp_no)
+  | EVAL is_expected = count >= 2 AND count <= 48 AND
                        avg_emp_no > 10055 AND avg_emp_no < 10095
   | KEEP is_expected
 ;
@@ -71,14 +69,13 @@ true
 
 
 after where
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | WHERE emp_no <= 10050 
   | SAMPLE 0.5
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no)), avg_emp_no = AVG(emp_no)
-  | EVAL is_expected = count >= 5 AND count <= 95 AND
-                       values_count >= 2 AND values_count <= 48 AND
+  | STATS count = COUNT(), avg_emp_no = AVG(emp_no)
+  | EVAL is_expected = count >= 2 AND count <= 48 AND
                        avg_emp_no > 10005 AND avg_emp_no < 10045
   | KEEP is_expected
 ;
@@ -89,14 +86,13 @@ true
 
 
 before sort
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SAMPLE 0.5
   | SORT emp_no
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no)), avg_emp_no = AVG(emp_no)
-  | EVAL is_expected = count >= 20 AND count <= 180 AND
-                       values_count >= 10 AND values_count <= 90 AND
+  | STATS count = COUNT(), avg_emp_no = AVG(emp_no)
+  | EVAL is_expected = count >= 10 AND count <= 90 AND
                        avg_emp_no > 10010 AND avg_emp_no < 10090
   | KEEP is_expected
 ;
@@ -107,14 +103,13 @@ true
 
 
 after sort
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SORT emp_no
   | SAMPLE 0.5
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no)), avg_emp_no = AVG(emp_no)
-  | EVAL is_expected = count >= 20 AND count <= 180 AND
-                       values_count >= 10 AND values_count <= 90 AND
+  | STATS count = COUNT(), avg_emp_no = AVG(emp_no)
+  | EVAL is_expected = count >= 10 AND count <= 90 AND
                        avg_emp_no > 10010 AND avg_emp_no < 10090
   | KEEP is_expected
 ;
@@ -125,13 +120,13 @@ true
 
 
 before limit
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SAMPLE 0.5
   | LIMIT 10
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no))
-  | EVAL is_expected = count == 10 AND values_count == 10
+  | STATS count = COUNT(emp_no)
+  | EVAL is_expected = count == 10
   | KEEP is_expected
 ;
 
@@ -141,14 +136,13 @@ true
 
 
 after limit
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | LIMIT 50
   | SAMPLE 0.5
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no))
-  | EVAL is_expected = count >= 5 AND count <= 95 AND
-                       values_count >= 2 AND values_count <= 48
+  | STATS count = COUNT(emp_no)
+  | EVAL is_expected = count >= 2 AND count <= 48
   | KEEP is_expected
 ;
 
@@ -158,7 +152,7 @@ true
 
 
 before mv_expand
-required_capability: sample_v2
+required_capability: sample_v3
 
 ROW x = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50], y = [1,2]
   | MV_EXPAND x
@@ -176,7 +170,7 @@ true
 
 
 after mv_expand
-required_capability: sample_v2
+required_capability: sample_v3
 
 ROW x = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50], y = [1,2]
   | MV_EXPAND x
@@ -194,15 +188,14 @@ true
 
 
 multiple samples
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SAMPLE 0.7
   | SAMPLE 0.8
   | SAMPLE 0.9
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(emp_no)), avg_emp_no = AVG(emp_no)
-  | EVAL is_expected = count >= 20 AND count <= 180 AND
-                       values_count >= 10 AND values_count <= 90 AND
+  | STATS count = COUNT(), avg_emp_no = AVG(emp_no)
+  | EVAL is_expected = count >= 10 AND count <= 90 AND
                        avg_emp_no > 10010 AND avg_emp_no < 10090
   | KEEP is_expected
 ;
@@ -213,15 +206,14 @@ true
 
 
 after stats
-required_capability: sample_v2
+required_capability: sample_v3
 
 FROM employees
   | SAMPLE 0.5
   | STATS avg_salary = AVG(salary) BY job_positions
   | SAMPLE 0.8
-  | STATS count = COUNT(), values_count = MV_COUNT(VALUES(avg_salary)), avg_avg_salary = AVG(avg_salary)
-  | EVAL is_expected = count >= 1 AND count <= 20 AND
-                       values_count >= 1 AND values_count <= 16 AND
+  | STATS count = COUNT(), avg_avg_salary = AVG(avg_salary)
+  | EVAL is_expected = count >= 1 AND count <= 16 AND
                        avg_avg_salary > 25000 AND avg_avg_salary < 75000
   | KEEP is_expected
 ;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1082,7 +1082,7 @@ public enum Cap {
         /**
          * Support for the SAMPLE command
          */
-        SAMPLE_V2(Build.current().isSnapshot()),
+        SAMPLE_V3(Build.current().isSnapshot()),
 
         /**
          * The {@code _query} API now gives a cast recommendation if multiple types are found in certain instances.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Count.java

@@ -25,10 +25,8 @@
 import org.elasticsearch.xpack.esql.expression.function.FunctionType;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromAggregateMetricDouble;
-import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvCount;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
 import org.elasticsearch.xpack.esql.planner.ToAggregator;
 
@@ -39,11 +37,9 @@
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 
-public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression, HasSampleCorrection {
+public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Count", Count::new);
 
-    private final boolean isSampleCorrected;
-
     @FunctionInfo(
         returnType = "long",
         description = "Returns the total number (count) of input values.",
@@ -98,20 +94,11 @@ public Count(
     }
 
     public Count(Source source, Expression field, Expression filter) {
-        this(source, field, filter, false);
-    }
-
-    private Count(Source source, Expression field, Expression filter, boolean isSampleCorrected) {
         super(source, field, filter, emptyList());
-        this.isSampleCorrected = isSampleCorrected;
     }
 
     private Count(StreamInput in) throws IOException {
         super(in);
-        // isSampleCorrected is only used during query optimization to mark
-        // whether this function has been processed. Hence there's no need to
-        // serialize it.
-        this.isSampleCorrected = false;
     }
 
     @Override
@@ -182,14 +169,4 @@ public Expression surrogate() {
 
         return null;
     }
-
-    @Override
-    public boolean isSampleCorrected() {
-        return isSampleCorrected;
-    }
-
-    @Override
-    public Expression sampleCorrection(Expression sampleProbability) {
-        return new ToLong(source(), new Div(source(), new Count(source(), field(), filter(), true), sampleProbability));
-    }
 }