Sampling prototype

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

@@ -331,6 +331,7 @@ static TransportVersion def(int id) {
  public static final TransportVersion ESQL_LOOKUP_JOIN_ON_EXPRESSION = def(9_163_0_00);
  public static final TransportVersion INFERENCE_REQUEST_ADAPTIVE_RATE_LIMITING_REMOVED = def(9_164_0_00);
  public static final TransportVersion SEARCH_SOURCE_EXCLUDE_INFERENCE_FIELDS_PARAM = def(9_165_0_00);
+ public static final TransportVersion RANDOM_SAMPLING = def(9_166_0_00);
 
  /*
  * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -405,6 +405,15 @@
 import org.elasticsearch.rest.action.synonyms.RestGetSynonymsSetsAction;
 import org.elasticsearch.rest.action.synonyms.RestPutSynonymRuleAction;
 import org.elasticsearch.rest.action.synonyms.RestPutSynonymsAction;
+import org.elasticsearch.sample.GetSampleAction;
+import org.elasticsearch.sample.GetSampleStatsAction;
+import org.elasticsearch.sample.PutSampleConfigAction;
+import org.elasticsearch.sample.RestGetSampleAction;
+import org.elasticsearch.sample.RestGetSampleStatsAction;
+import org.elasticsearch.sample.RestPutSampleConfigAction;
+import org.elasticsearch.sample.TransportGetSampleAction;
+import org.elasticsearch.sample.TransportGetSampleStatsAction;
+import org.elasticsearch.sample.TransportPutSampleConfigAction;
 import org.elasticsearch.snapshots.TransportUpdateSnapshotStatusAction;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.telemetry.TelemetryProvider;
@@ -813,6 +822,10 @@ public <Request extends ActionRequest, Response extends ActionResponse> void reg
  actions.register(GetSynonymRuleAction.INSTANCE, TransportGetSynonymRuleAction.class);
  actions.register(DeleteSynonymRuleAction.INSTANCE, TransportDeleteSynonymRuleAction.class);
 
+ actions.register(PutSampleConfigAction.INSTANCE, TransportPutSampleConfigAction.class);
+ actions.register(GetSampleAction.INSTANCE, TransportGetSampleAction.class);
+ actions.register(GetSampleStatsAction.INSTANCE, TransportGetSampleStatsAction.class);
+
  return unmodifiableMap(actions.getRegistry());
  }
 
@@ -1040,6 +1053,10 @@ public void initRestHandlers(Supplier<DiscoveryNodes> nodesInCluster, Predicate<
  registerHandler.accept(new RestPutSynonymRuleAction());
  registerHandler.accept(new RestGetSynonymRuleAction());
  registerHandler.accept(new RestDeleteSynonymRuleAction());
+
+ registerHandler.accept(new RestPutSampleConfigAction());
+ registerHandler.accept(new RestGetSampleAction(projectIdResolver));
+ registerHandler.accept(new RestGetSampleStatsAction(projectIdResolver));
  }
 
  @Override

server/src/main/java/org/elasticsearch/action/bulk/TransportAbstractBulkAction.java

@@ -64,7 +64,6 @@
  */
 public abstract class TransportAbstractBulkAction extends HandledTransportAction<BulkRequest, BulkResponse> {
  private static final Logger logger = LogManager.getLogger(TransportAbstractBulkAction.class);
-
  public static final Set<String> STREAMS_ALLOWED_PARAMS = new HashSet<>(9) {
  {
  add("error_trace");

server/src/main/java/org/elasticsearch/cluster/ClusterModule.java

@@ -90,6 +90,7 @@
 import org.elasticsearch.persistent.PersistentTasksCustomMetadata;
 import org.elasticsearch.persistent.PersistentTasksNodeService;
 import org.elasticsearch.plugins.ClusterPlugin;
+import org.elasticsearch.sample.TransportPutSampleConfigAction;
 import org.elasticsearch.script.ScriptMetadata;
 import org.elasticsearch.snapshots.RegisteredPolicySnapshots;
 import org.elasticsearch.snapshots.SnapshotsInfoService;
@@ -277,6 +278,12 @@ public static List<Entry> getNamedWriteables() {
  PersistentTasksCustomMetadata::new,
  PersistentTasksCustomMetadata::readDiffFrom
  );
+ registerProjectCustom(
+ entries,
+ TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME,
+ TransportPutSampleConfigAction.SamplingConfigCustomMetadata::new,
+ TransportPutSampleConfigAction.SamplingConfigCustomMetadata::readDiffFrom
+ );
  // Cluster scoped persistent tasks
  registerMetadataCustom(
  entries,
@@ -357,6 +364,13 @@ public static List<NamedXContentRegistry.Entry> getNamedXWriteables() {
  ClusterPersistentTasksCustomMetadata::fromXContent
  )
  );
+ entries.add(
+ new NamedXContentRegistry.Entry(
+ Metadata.ProjectCustom.class,
+ new ParseField(TransportPutSampleConfigAction.SamplingConfigCustomMetadata.NAME),
+ TransportPutSampleConfigAction.SamplingConfigCustomMetadata::fromXContent
+ )
+ );
  entries.add(
  new NamedXContentRegistry.Entry(
  Metadata.ProjectCustom.class,

server/src/main/java/org/elasticsearch/ingest/IngestService.java

@@ -52,7 +52,6 @@
 import org.elasticsearch.cluster.service.MasterServiceTaskQueue;
 import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.TriConsumer;
-import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.collect.ImmutableOpenMap;
 import org.elasticsearch.common.logging.DeprecationCategory;
 import org.elasticsearch.common.logging.DeprecationLogger;
@@ -1390,45 +1389,10 @@ private void attemptToSampleData(
  * We need both the original document and the fully updated document for sampling, so we make a copy of the original
  * before overwriting it here. We can discard it after sampling.
  */
- samplingService.maybeSample(projectMetadata, indexRequest.index(), () -> {
- IndexRequest original = copyIndexRequestForSampling(indexRequest);
- updateIndexRequestMetadata(original, originalDocumentMetadata);
- return original;
- }, ingestDocument);
-
+ samplingService.maybeSample(projectMetadata, originalDocumentMetadata.getIndex(), indexRequest, ingestDocument);
  }
  }
 
- /**
- * Creates a copy of an IndexRequest to be used by random sampling.
- * @param original The IndexRequest to be copied
- * @return A copy of the IndexRequest
- */
- private IndexRequest copyIndexRequestForSampling(IndexRequest original) {
- IndexRequest clonedRequest = new IndexRequest(original.index());
- clonedRequest.id(original.id());
- clonedRequest.routing(original.routing());
- clonedRequest.version(original.version());
- clonedRequest.versionType(original.versionType());
- clonedRequest.setPipeline(original.getPipeline());
- clonedRequest.setFinalPipeline(original.getFinalPipeline());
- clonedRequest.setIfSeqNo(original.ifSeqNo());
- clonedRequest.setIfPrimaryTerm(original.ifPrimaryTerm());
- clonedRequest.setRefreshPolicy(original.getRefreshPolicy());
- clonedRequest.waitForActiveShards(original.waitForActiveShards());
- clonedRequest.timeout(original.timeout());
- clonedRequest.opType(original.opType());
- clonedRequest.setParentTask(original.getParentTask());
- clonedRequest.setRequireDataStream(original.isRequireDataStream());
- clonedRequest.setRequireAlias(original.isRequireAlias());
- clonedRequest.setIncludeSourceOnError(original.getIncludeSourceOnError());
- BytesReference source = original.source();
- if (source != null) {
- clonedRequest.source(source, original.getContentType());
- }
- return clonedRequest;
- }
-
  private static void executePipeline(
  final IngestDocument ingestDocument,
  final Pipeline pipeline,