Code review comment: Correct, thanks @hermes85pl  for pointing out, it should not be done for every element processed

Code under review:
public boolean start() throws IOException {
      restClient = source.spec.getConnectionConfiguration().createClient();

      String query = source.spec.getQuery() != null ? source.spec.getQuery().get() : null;
      if (query == null) {
        query = "{\"query\": { \"match_all\": {} }}";
      }
      if ((source.backendVersion >= 5) && source.numSlices != null && source.numSlices > 1) {
        
        String sliceQuery =
            String.format("\"slice\": {\"id\": %s,\"max\": %s}", source.sliceId, source.numSlices);
        query = query.replaceFirst("\\{", "{" + sliceQuery + ",");
      }
      String endPoint =
          String.format(
              "/%s/%s/_search",
              source.spec.getConnectionConfiguration().getIndex(),
              source.spec.getConnectionConfiguration().getType());
      Map<String, String> params = new HashMap<>();
      params.put("scroll", source.spec.getScrollKeepalive());
      if (source.backendVersion == 2) {
        params.put("size", String.valueOf(source.spec.getBatchSize()));
        if (source.shardPreference != null) {
          params.put("preference", "_shards:" + source.shardPreference);
        }
      }
      HttpEntity queryEntity = new NStringEntity(query, ContentType.APPLICATION_JSON);
      Request request = new Request("GET", endPoint);
      request.addParameters(params);
      request.setEntity(queryEntity);
      Response response = restClient.performRequest(request);
      JsonNode searchResult = parseResponse(response.getEntity());
      updateScrollId(searchResult);
      return readNextBatchAndReturnFirstDocument(searchResult);
    }

    private void updateScrollId(JsonNode searchResult) {
      scrollId = searchResult.path("_scroll_id").asText();
    }

    @Override
    public boolean advance() throws IOException {
      if (batchIterator.hasNext()) {
        current = batchIterator.next();
        return true;
      } else {
        String requestBody =
            String.format(
                "{\"scroll\" : \"%s\",\"scroll_id\" : \"%s\"}",
                source.spec.getScrollKeepalive(), scrollId);
        HttpEntity scrollEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
        Request request = new Request("GET", "/_search/scroll");
        request.addParameters(Collections.emptyMap());
        request.setEntity(scrollEntity);
        Response response = restClient.performRequest(request);
        JsonNode searchResult = parseResponse(response.getEntity());
        updateScrollId(searchResult);
        return readNextBatchAndReturnFirstDocument(searchResult);
      }
    }

    private boolean readNextBatchAndReturnFirstDocument(JsonNode searchResult) {
      
      JsonNode hits = searchResult.path("hits").path("hits");
      if (hits.size() == 0) {
        current = null;
        batchIterator = null;
        return false;
      }
      
      List<String> batch = new ArrayList<>();
      boolean withMetadata = source.spec.isWithMetadata();
      for (JsonNode hit : hits) {
        if (withMetadata) {
          batch.add(hit.toString());
        } else {
          String document = hit.path("_source").toString();
          batch.add(document);
        }
      }
      batchIterator = batch.listIterator();
      current = batchIterator.next();
      return true;
    }

    @Override
    public String getCurrent() throws NoSuchElementException {
      if (current == null) {
        throw new NoSuchElementException();
      }
      return current;
    }

    @Override
    public void close() throws IOException {
      
      String requestBody = String.format("{\"scroll_id\" : [\"%s\"]}", scrollId);
      HttpEntity entity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
      try {
        Request request = new Request("DELETE", "/_search/scroll");
        request.addParameters(Collections.emptyMap());
        request.setEntity(entity);
        restClient.performRequest(request);
      } finally {
        if (restClient != null) {
          restClient.close();
        }
      }
    }

    @Override
    public BoundedSource<String> getCurrentSource() {
      return source;
    }
  }
  /**
   * A POJO encapsulating a configuration for retry behavior when issuing requests to ES. A retry
   * will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes first, for
   * 429 TOO_MANY_REQUESTS error.
   */
  @AutoValue
  public abstract static class RetryConfiguration implements Serializable {
    @VisibleForTesting
    static final RetryPredicate DEFAULT_RETRY_PREDICATE = new DefaultRetryPredicate();

    abstract int getMaxAttempts();

    abstract Duration getMaxDuration();

    abstract RetryPredicate getRetryPredicate();

    abstract Builder builder();

    @AutoValue.Builder
    abstract static class Builder {
      abstract ElasticsearchIO.RetryConfiguration.Builder setMaxAttempts(int maxAttempts);

      abstract ElasticsearchIO.RetryConfiguration.Builder setMaxDuration(Duration maxDuration);

      abstract ElasticsearchIO.RetryConfiguration.Builder setRetryPredicate(
          RetryPredicate retryPredicate);

      abstract ElasticsearchIO.RetryConfiguration build();
    }

    /**
     * Creates RetryConfiguration for {@link ElasticsearchIO} with provided maxAttempts,
     * maxDurations and exponential backoff based retries.
     *
     * @param maxAttempts max number of attempts.
     * @param maxDuration maximum duration for retries.
     * @return {@link RetryConfiguration} object with provided settings.
     */
    public static RetryConfiguration create(int maxAttempts, Duration maxDuration) {
      checkArgument(maxAttempts > 0, "maxAttempts must be greater than 0");
      checkArgument(
          maxDuration != null && maxDuration.isLongerThan(Duration.ZERO),
          "maxDuration must be greater than 0");
      return new AutoValue_ElasticsearchIO_RetryConfiguration.Builder()
          .setMaxAttempts(maxAttempts)
          .setMaxDuration(maxDuration)
          .setRetryPredicate(DEFAULT_RETRY_PREDICATE)
          .build();
    }

    
    @VisibleForTesting
    RetryConfiguration withRetryPredicate(RetryPredicate predicate) {
      checkArgument(predicate != null, "predicate must be provided");
      return builder().setRetryPredicate(predicate).build();
    }

    /**
     * An interface used to control if we retry the Elasticsearch call when a {@link Response} is
     * obtained. If {@link RetryPredicate
     * the requests to the Elasticsearch server if the {@link RetryConfiguration} permits it.
     */
    @FunctionalInterface
    interface RetryPredicate extends Predicate<HttpEntity>, Serializable {}

    /**
     * This is the default predicate used to test if a failed ES operation should be retried. A
     * retry will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes
     * first, for TOO_MANY_REQUESTS(429) error.
     */
    @VisibleForTesting
    static class DefaultRetryPredicate implements RetryPredicate {

      private int errorCode;

      DefaultRetryPredicate(int code) {
        this.errorCode = code;
      }

      DefaultRetryPredicate() {
        this(429);
      }

      /** Returns true if the response has the error code for any mutation. */
      private static boolean errorCodePresent(HttpEntity responseEntity, int errorCode) {
        try {
          JsonNode json = parseResponse(responseEntity);
          if (json.path("errors").asBoolean()) {
            for (JsonNode item : json.path("items")) {
              if (item.findValue("status").asInt() == errorCode) {
                return true;
              }
            }
          }
        } catch (IOException e) {
          LOG.warn("Could not extract error codes from responseEntity {}", responseEntity);
        }
        return false;
      }

      @Override
      public boolean test(HttpEntity responseEntity) {
        return errorCodePresent(responseEntity, errorCode);
      }
    }
  }

  /** A {@link PTransform} writing data to Elasticsearch. */
  @AutoValue
  public abstract static class Write extends PTransform<PCollection<String>, PDone> {

    /**
     * Interface allowing a specific field value to be returned from a parsed JSON document. This is
     * used for using explicit document ids, and for dynamic routing (index/Type) on a document
     * basis. A null response will result in default behaviour and an exception will be propagated
     * as a failure.
     */
    public interface FieldValueExtractFn extends SerializableFunction<JsonNode, String> {}

    @Nullable
    abstract ConnectionConfiguration getConnectionConfiguration();

    abstract long getMaxBatchSize();

    abstract long getMaxBatchSizeBytes();

    @Nullable
    abstract FieldValueExtractFn getIdFn();

    @Nullable
    abstract FieldValueExtractFn getIndexFn();

    @Nullable
    abstract FieldValueExtractFn getTypeFn();

    @Nullable
    abstract RetryConfiguration getRetryConfiguration();

    abstract boolean getUsePartialUpdate();

    abstract Builder builder();

    @AutoValue.Builder
    abstract static class Builder {
      abstract Builder setConnectionConfiguration(ConnectionConfiguration connectionConfiguration);

      abstract Builder setMaxBatchSize(long maxBatchSize);

      abstract Builder setMaxBatchSizeBytes(long maxBatchSizeBytes);

      abstract Builder setIdFn(FieldValueExtractFn idFunction);

      abstract Builder setIndexFn(FieldValueExtractFn indexFn);

      abstract Builder setTypeFn(FieldValueExtractFn typeFn);

      abstract Builder setUsePartialUpdate(boolean usePartialUpdate);

      abstract Builder setRetryConfiguration(RetryConfiguration retryConfiguration);

      abstract Write build();
    }

    /**
     * Provide the Elasticsearch connection configuration object.
     *
     * @param connectionConfiguration the Elasticsearch {@link ConnectionConfiguration} object
     * @return the {@link Write} with connection configuration set
     */
    public Write withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
      checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
      return builder().setConnectionConfiguration(connectionConfiguration).build();
    }

    /**
     * Provide a maximum size in number of documents for the batch see bulk API
     * (https:
     * docs (like Elasticsearch bulk size advice). See
     * https:
     * execution engine, size of bundles may vary, this sets the maximum size. Change this if you
     * need to have smaller Elasticsearch bulks.
     *
     * @param batchSize maximum batch size in number of documents
     * @return the {@link Write} with connection batch size set
     */
    public Write withMaxBatchSize(long batchSize) {
      checkArgument(batchSize > 0, "batchSize must be > 0, but was %s", batchSize);
      return builder().setMaxBatchSize(batchSize).build();
    }

    /**
     * Provide a maximum size in bytes for the batch see bulk API
     * (https:
     * (like Elasticsearch bulk size advice). See
     * https:
     * execution engine, size of bundles may vary, this sets the maximum size. Change this if you
     * need to have smaller Elasticsearch bulks.
     *
     * @param batchSizeBytes maximum batch size in bytes
     * @return the {@link Write} with connection batch size in bytes set
     */
    public Write withMaxBatchSizeBytes(long batchSizeBytes) {
      checkArgument(batchSizeBytes > 0, "batchSizeBytes must be > 0, but was %s", batchSizeBytes);
      return builder().setMaxBatchSizeBytes(batchSizeBytes).build();
    }

    /**
     * Provide a function to extract the id from the document. This id will be used as the document
     * id in Elasticsearch. Should the function throw an Exception then the batch will fail and the
     * exception propagated.
     *
     * @param idFn to extract the document ID
     * @return the {@link Write} with the function set
     */
    public Write withIdFn(FieldValueExtractFn idFn) {
      checkArgument(idFn != null, "idFn must not be null");
      return builder().setIdFn(idFn).build();
    }

    /**
     * Provide a function to extract the target index from the document allowing for dynamic
     * document routing. Should the function throw an Exception then the batch will fail and the
     * exception propagated.
     *
     * @param indexFn to extract the destination index from
     * @return the {@link Write} with the function set
     */
    public Write withIndexFn(FieldValueExtractFn indexFn) {
      checkArgument(indexFn != null, "indexFn must not be null");
      return builder().setIndexFn(indexFn).build();
    }

    /**
     * Provide a function to extract the target type from the document allowing for dynamic document
     * routing. Should the function throw an Exception then the batch will fail and the exception
     * propagated. Users are encouraged to consider carefully if multipe types are a sensible model
     * <a
     * href="https:
     * discussed in this blog</a>.
     *
     * @param typeFn to extract the destination index from
     * @return the {@link Write} with the function set
     */
    public Write withTypeFn(FieldValueExtractFn typeFn) {
      checkArgument(typeFn != null, "typeFn must not be null");
      return builder().setTypeFn(typeFn).build();
    }

    /**
     * Provide an instruction to control whether partial updates or inserts (default) are issued to
     * Elasticsearch.
     *
     * @param usePartialUpdate set to true to issue partial updates
     * @return the {@link Write} with the partial update control set
     */
    public Write withUsePartialUpdate(boolean usePartialUpdate) {
      return builder().setUsePartialUpdate(usePartialUpdate).build();
    }

    /**
     * Provides configuration to retry a failed batch call to Elasticsearch. A batch is considered
     * as failed if the underlying {@link RestClient} surfaces 429 HTTP status code as error for one
     * or more of the items in the {@link Response}. Users should consider that retrying might
     * compound the underlying problem which caused the initial failure. Users should also be aware
     * that once retrying is exhausted the error is surfaced to the runner which <em>may</em> then
     * opt to retry the current bundle in entirety or abort if the max number of retries of the
     * runner is completed. Retrying uses an exponential backoff algorithm, with minimum backoff of
     * 5 seconds and then surfacing the error once the maximum number of retries or maximum
     * configuration duration is exceeded.
     *
     * <p>Example use:
     *
     * <pre>{@code
     * ElasticsearchIO.write()
     *   .withRetryConfiguration(ElasticsearchIO.RetryConfiguration.create(10, Duration.standardMinutes(3))
     *   ...
     * }</pre>
     *
     * @param retryConfiguration the rules which govern the retry behavior
     * @return the {@link Write} with retrying configured
     */
    public Write withRetryConfiguration(RetryConfiguration retryConfiguration) {
      checkArgument(retryConfiguration != null, "retryConfiguration is required");
      return builder().setRetryConfiguration(retryConfiguration).build();
    }

    @Override
    public PDone expand(PCollection<String> input) {
      ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
      checkState(connectionConfiguration != null, "withConnectionConfiguration() is required");
      input.apply(ParDo.of(new WriteFn(this)));
      return PDone.in(input.getPipeline());
    }

    /** {@link DoFn} to for the {@link Write} transform. */
    @VisibleForTesting
    static class WriteFn extends DoFn<String, Void> {
      private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
      private static final int DEFAULT_RETRY_ON_CONFLICT = 5; 

      private static final Duration RETRY_INITIAL_BACKOFF = Duration.standardSeconds(5);

      @VisibleForTesting
      static final String RETRY_ATTEMPT_LOG = "Error writing to Elasticsearch. Retry attempt[%d]";

      @VisibleForTesting
      static final String RETRY_FAILED_LOG =
          "Error writing to ES after %d attempt(s). No more attempts allowed";

      private transient FluentBackoff retryBackoff;

      private int backendVersion;
      private final Write spec;
      private transient RestClient restClient;
      private ArrayList<String> batch;
      private long currentBatchSizeBytes;

      
      private static class DocumentMetadata implements Serializable {
        final String index;
        final String type;
        final String id;
        final Integer retryOnConflict;

        DocumentMetadata(String index, String type, String id, Integer retryOnConflict) {
          this.index = index;
          this.type = type;
          this.id = id;
          this.retryOnConflict = retryOnConflict;
        }
      }

      @VisibleForTesting
      WriteFn(Write spec) {
        this.spec = spec;
      }

      @Setup
      public void setup() throws IOException {
        ConnectionConfiguration connectionConfiguration = spec.getConnectionConfiguration();
        backendVersion = getBackendVersion(connectionConfiguration);
        restClient = connectionConfiguration.createClient();

        retryBackoff =
            FluentBackoff.DEFAULT.withMaxRetries(0).withInitialBackoff(RETRY_INITIAL_BACKOFF);

        if (spec.getRetryConfiguration() != null) {
          retryBackoff =
              FluentBackoff.DEFAULT
                  .withInitialBackoff(RETRY_INITIAL_BACKOFF)
                  .withMaxRetries(spec.getRetryConfiguration().getMaxAttempts() - 1)
                  .withMaxCumulativeBackoff(spec.getRetryConfiguration().getMaxDuration());
        }
      }

      @StartBundle
      public void startBundle(StartBundleContext context) {
        batch = new ArrayList<>();
        currentBatchSizeBytes = 0;
      }

      private class DocumentMetadataSerializer extends StdSerializer<DocumentMetadata> {

        private DocumentMetadataSerializer() {
          super(DocumentMetadata.class);
        }

        @Override
        public void serialize(
            DocumentMetadata value, JsonGenerator gen, SerializerProvider provider)
            throws IOException {
          gen.writeStartObject();
          if (value.index != null) {
            gen.writeStringField("_index", value.index);
          }
          if (value.type != null) {
            gen.writeStringField("_type", value.type);
          }
          if (value.id != null) {
            gen.writeStringField("_id", value.id);
          }
          if (value.retryOnConflict != null && (backendVersion <= 6)) {
            gen.writeNumberField("_retry_on_conflict", value.retryOnConflict);
          }
          if (value.retryOnConflict != null && backendVersion >= 7) {
            gen.writeNumberField("retry_on_conflict", value.retryOnConflict);
          }
          gen.writeEndObject();
        }
      }
      /**
       * Extracts the components that comprise the document address from the document using the
       * {@link FieldValueExtractFn} configured. This allows any or all of the index, type and
       * document id to be controlled on a per document basis. If none are provided then an empty
       * default of {@code {}} is returned. Sanitization of the index is performed, automatically
       * lower-casing the value as required by Elasticsearch.
       *
       * @param document the json from which the index, type and id may be extracted
       * @return the document address as JSON or the default
       * @throws IOException if the document cannot be parsed as JSON
       */
      private String getDocumentMetadata(String document) throws IOException {
        if (spec.getIndexFn() != null || spec.getTypeFn() != null || spec.getIdFn() != null) {
          
          JsonNode parsedDocument = OBJECT_MAPPER.readTree(document);

          DocumentMetadata metadata =
              new DocumentMetadata(
                  spec.getIndexFn() != null
                      ? lowerCaseOrNull(spec.getIndexFn().apply(parsedDocument))
                      : null,
                  spec.getTypeFn() != null ? spec.getTypeFn().apply(parsedDocument) : null,
                  spec.getIdFn() != null ? spec.getIdFn().apply(parsedDocument) : null,
                  spec.getUsePartialUpdate() ? DEFAULT_RETRY_ON_CONFLICT : null);
          SimpleModule module = new SimpleModule();
          module.addSerializer(DocumentMetadata.class, new DocumentMetadataSerializer());
          OBJECT_MAPPER.registerModule(module);
          return OBJECT_MAPPER.writeValueAsString(metadata);
        } else {
          return "{}"; 
        }
      }

      private static String lowerCaseOrNull(String input) {
        return input == null ? null : input.toLowerCase();
      }

      @ProcessElement
      public void processElement(ProcessContext context) throws Exception {
        String document = context.element();
        String documentMetadata = getDocumentMetadata(document);

        
        if (spec.getUsePartialUpdate()) {
          batch.add(
              String.format(
                  "{ \"update\" : %s }%n{ \"doc\" : %s, \"doc_as_upsert\" : true }%n",
                  documentMetadata, document));
        } else {
          batch.add(String.format("{ \"index\" : %s }%n%s%n", documentMetadata, document));
        }

        currentBatchSizeBytes += document.getBytes(StandardCharsets.UTF_8).length;
        if (batch.size() >= spec.getMaxBatchSize()
            || currentBatchSizeBytes >= spec.getMaxBatchSizeBytes()) {
          flushBatch();
        }
      }

      @FinishBundle
      public void finishBundle(FinishBundleContext context)
          throws IOException, InterruptedException {
        flushBatch();
      }

      private void flushBatch() throws IOException, InterruptedException {
        if (batch.isEmpty()) {
          return;
        }
        StringBuilder bulkRequest = new StringBuilder();
        for (String json : batch) {
          bulkRequest.append(json);
        }
        batch.clear();
        currentBatchSizeBytes = 0;
        Response response;
        HttpEntity responseEntity;
        
        
        
        String endPoint =
            String.format(
                "/%s/%s/_bulk",
                spec.getConnectionConfiguration().getIndex(),
                spec.getConnectionConfiguration().getType());
        HttpEntity requestBody =
            new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
        Request request = new Request("POST", endPoint);
        request.addParameters(Collections.emptyMap());
        request.setEntity(requestBody);
        response = restClient.performRequest(request);
        responseEntity = new BufferedHttpEntity(response.getEntity());
        if (spec.getRetryConfiguration() != null
            && spec.getRetryConfiguration().getRetryPredicate().test(responseEntity)) {
          responseEntity = handleRetry("POST", endPoint, Collections.emptyMap(), requestBody);
        }
        checkForErrors(responseEntity, backendVersion, spec.getUsePartialUpdate());
      }

      /** retry request based on retry configuration policy. */
      private HttpEntity handleRetry(
          String method, String endpoint, Map<String, String> params, HttpEntity requestBody)
          throws IOException, InterruptedException {
        Response response;
        HttpEntity responseEntity;
        Sleeper sleeper = Sleeper.DEFAULT;
        BackOff backoff = retryBackoff.backoff();
        int attempt = 0;
        
        while (BackOffUtils.next(sleeper, backoff)) {
          LOG.warn(String.format(RETRY_ATTEMPT_LOG, ++attempt));
          Request request = new Request(method, endpoint);
          request.addParameters(params);
          request.setEntity(requestBody);
          response = restClient.performRequest(request);
          responseEntity = new BufferedHttpEntity(response.getEntity());
          
          if (!spec.getRetryConfiguration().getRetryPredicate().test(responseEntity)) {
            return responseEntity;
          }
        }
        throw new IOException(String.format(RETRY_FAILED_LOG, attempt));
      }

      @Teardown
      public void closeClient() throws IOException {
        if (restClient != null) {
          restClient.close();
        }
      }
    }
  }

  static int getBackendVersion(ConnectionConfiguration connectionConfiguration) {
    try (RestClient restClient = connectionConfiguration.createClient()) {
      Request request = new Request("GET", "");
      Response response = restClient.performRequest(request);
      JsonNode jsonNode = parseResponse(response.getEntity());
      int backendVersion =
          Integer.parseInt(jsonNode.path("version").path("number").asText().substring(0, 1));
      checkArgument(
          (backendVersion == 2
              || backendVersion == 5
              || backendVersion == 6
              || backendVersion == 7),
          "The Elasticsearch version to connect to is %s.x. "
              + "This version of the ElasticsearchIO is only compatible with "
              + "Elasticsearch v7.x, v6.x, v5.x and v2.x",
          backendVersion);
      return backendVersion;

    } catch (IOException e) {
      throw new IllegalArgumentException("Cannot get Elasticsearch version", e);
    }
  }
}

Line commented on: OBJECT_MAPPER.registerModule(module);

Revised code:
public boolean start() throws IOException {
      restClient = source.spec.getConnectionConfiguration().createClient();

      String query = source.spec.getQuery() != null ? source.spec.getQuery().get() : null;
      if (query == null) {
        query = "{\"query\": { \"match_all\": {} }}";
      }
      if ((source.backendVersion >= 5) && source.numSlices != null && source.numSlices > 1) {
        
        String sliceQuery =
            String.format("\"slice\": {\"id\": %s,\"max\": %s}", source.sliceId, source.numSlices);
        query = query.replaceFirst("\\{", "{" + sliceQuery + ",");
      }
      String endPoint =
          String.format(
              "/%s/%s/_search",
              source.spec.getConnectionConfiguration().getIndex(),
              source.spec.getConnectionConfiguration().getType());
      Map<String, String> params = new HashMap<>();
      params.put("scroll", source.spec.getScrollKeepalive());
      if (source.backendVersion == 2) {
        params.put("size", String.valueOf(source.spec.getBatchSize()));
        if (source.shardPreference != null) {
          params.put("preference", "_shards:" + source.shardPreference);
        }
      }
      HttpEntity queryEntity = new NStringEntity(query, ContentType.APPLICATION_JSON);
      Request request = new Request("GET", endPoint);
      request.addParameters(params);
      request.setEntity(queryEntity);
      Response response = restClient.performRequest(request);
      JsonNode searchResult = parseResponse(response.getEntity());
      updateScrollId(searchResult);
      return readNextBatchAndReturnFirstDocument(searchResult);
    }

    private void updateScrollId(JsonNode searchResult) {
      scrollId = searchResult.path("_scroll_id").asText();
    }

    @Override
    public boolean advance() throws IOException {
      if (batchIterator.hasNext()) {
        current = batchIterator.next();
        return true;
      } else {
        String requestBody =
            String.format(
                "{\"scroll\" : \"%s\",\"scroll_id\" : \"%s\"}",
                source.spec.getScrollKeepalive(), scrollId);
        HttpEntity scrollEntity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
        Request request = new Request("GET", "/_search/scroll");
        request.addParameters(Collections.emptyMap());
        request.setEntity(scrollEntity);
        Response response = restClient.performRequest(request);
        JsonNode searchResult = parseResponse(response.getEntity());
        updateScrollId(searchResult);
        return readNextBatchAndReturnFirstDocument(searchResult);
      }
    }

    private boolean readNextBatchAndReturnFirstDocument(JsonNode searchResult) {
      
      JsonNode hits = searchResult.path("hits").path("hits");
      if (hits.size() == 0) {
        current = null;
        batchIterator = null;
        return false;
      }
      
      List<String> batch = new ArrayList<>();
      boolean withMetadata = source.spec.isWithMetadata();
      for (JsonNode hit : hits) {
        if (withMetadata) {
          batch.add(hit.toString());
        } else {
          String document = hit.path("_source").toString();
          batch.add(document);
        }
      }
      batchIterator = batch.listIterator();
      current = batchIterator.next();
      return true;
    }

    @Override
    public String getCurrent() throws NoSuchElementException {
      if (current == null) {
        throw new NoSuchElementException();
      }
      return current;
    }

    @Override
    public void close() throws IOException {
      
      String requestBody = String.format("{\"scroll_id\" : [\"%s\"]}", scrollId);
      HttpEntity entity = new NStringEntity(requestBody, ContentType.APPLICATION_JSON);
      try {
        Request request = new Request("DELETE", "/_search/scroll");
        request.addParameters(Collections.emptyMap());
        request.setEntity(entity);
        restClient.performRequest(request);
      } finally {
        if (restClient != null) {
          restClient.close();
        }
      }
    }

    @Override
    public BoundedSource<String> getCurrentSource() {
      return source;
    }
  }
  /**
   * A POJO encapsulating a configuration for retry behavior when issuing requests to ES. A retry
   * will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes first, for
   * 429 TOO_MANY_REQUESTS error.
   */
  @AutoValue
  public abstract static class RetryConfiguration implements Serializable {
    @VisibleForTesting
    static final RetryPredicate DEFAULT_RETRY_PREDICATE = new DefaultRetryPredicate();

    abstract int getMaxAttempts();

    abstract Duration getMaxDuration();

    abstract RetryPredicate getRetryPredicate();

    abstract Builder builder();

    @AutoValue.Builder
    abstract static class Builder {
      abstract ElasticsearchIO.RetryConfiguration.Builder setMaxAttempts(int maxAttempts);

      abstract ElasticsearchIO.RetryConfiguration.Builder setMaxDuration(Duration maxDuration);

      abstract ElasticsearchIO.RetryConfiguration.Builder setRetryPredicate(
          RetryPredicate retryPredicate);

      abstract ElasticsearchIO.RetryConfiguration build();
    }

    /**
     * Creates RetryConfiguration for {@link ElasticsearchIO} with provided maxAttempts,
     * maxDurations and exponential backoff based retries.
     *
     * @param maxAttempts max number of attempts.
     * @param maxDuration maximum duration for retries.
     * @return {@link RetryConfiguration} object with provided settings.
     */
    public static RetryConfiguration create(int maxAttempts, Duration maxDuration) {
      checkArgument(maxAttempts > 0, "maxAttempts must be greater than 0");
      checkArgument(
          maxDuration != null && maxDuration.isLongerThan(Duration.ZERO),
          "maxDuration must be greater than 0");
      return new AutoValue_ElasticsearchIO_RetryConfiguration.Builder()
          .setMaxAttempts(maxAttempts)
          .setMaxDuration(maxDuration)
          .setRetryPredicate(DEFAULT_RETRY_PREDICATE)
          .build();
    }

    
    @VisibleForTesting
    RetryConfiguration withRetryPredicate(RetryPredicate predicate) {
      checkArgument(predicate != null, "predicate must be provided");
      return builder().setRetryPredicate(predicate).build();
    }

    /**
     * An interface used to control if we retry the Elasticsearch call when a {@link Response} is
     * obtained. If {@link RetryPredicate
     * the requests to the Elasticsearch server if the {@link RetryConfiguration} permits it.
     */
    @FunctionalInterface
    interface RetryPredicate extends Predicate<HttpEntity>, Serializable {}

    /**
     * This is the default predicate used to test if a failed ES operation should be retried. A
     * retry will be attempted until the maxAttempts or maxDuration is exceeded, whichever comes
     * first, for TOO_MANY_REQUESTS(429) error.
     */
    @VisibleForTesting
    static class DefaultRetryPredicate implements RetryPredicate {

      private int errorCode;

      DefaultRetryPredicate(int code) {
        this.errorCode = code;
      }

      DefaultRetryPredicate() {
        this(429);
      }

      /** Returns true if the response has the error code for any mutation. */
      private static boolean errorCodePresent(HttpEntity responseEntity, int errorCode) {
        try {
          JsonNode json = parseResponse(responseEntity);
          if (json.path("errors").asBoolean()) {
            for (JsonNode item : json.path("items")) {
              if (item.findValue("status").asInt() == errorCode) {
                return true;
              }
            }
          }
        } catch (IOException e) {
          LOG.warn("Could not extract error codes from responseEntity {}", responseEntity);
        }
        return false;
      }

      @Override
      public boolean test(HttpEntity responseEntity) {
        return errorCodePresent(responseEntity, errorCode);
      }
    }
  }

  /** A {@link PTransform} writing data to Elasticsearch. */
  @AutoValue
  public abstract static class Write extends PTransform<PCollection<String>, PDone> {

    /**
     * Interface allowing a specific field value to be returned from a parsed JSON document. This is
     * used for using explicit document ids, and for dynamic routing (index/Type) on a document
     * basis. A null response will result in default behaviour and an exception will be propagated
     * as a failure.
     */
    public interface FieldValueExtractFn extends SerializableFunction<JsonNode, String> {}

    @Nullable
    abstract ConnectionConfiguration getConnectionConfiguration();

    abstract long getMaxBatchSize();

    abstract long getMaxBatchSizeBytes();

    @Nullable
    abstract FieldValueExtractFn getIdFn();

    @Nullable
    abstract FieldValueExtractFn getIndexFn();

    @Nullable
    abstract FieldValueExtractFn getTypeFn();

    @Nullable
    abstract RetryConfiguration getRetryConfiguration();

    abstract boolean getUsePartialUpdate();

    abstract Builder builder();

    @AutoValue.Builder
    abstract static class Builder {
      abstract Builder setConnectionConfiguration(ConnectionConfiguration connectionConfiguration);

      abstract Builder setMaxBatchSize(long maxBatchSize);

      abstract Builder setMaxBatchSizeBytes(long maxBatchSizeBytes);

      abstract Builder setIdFn(FieldValueExtractFn idFunction);

      abstract Builder setIndexFn(FieldValueExtractFn indexFn);

      abstract Builder setTypeFn(FieldValueExtractFn typeFn);

      abstract Builder setUsePartialUpdate(boolean usePartialUpdate);

      abstract Builder setRetryConfiguration(RetryConfiguration retryConfiguration);

      abstract Write build();
    }

    /**
     * Provide the Elasticsearch connection configuration object.
     *
     * @param connectionConfiguration the Elasticsearch {@link ConnectionConfiguration} object
     * @return the {@link Write} with connection configuration set
     */
    public Write withConnectionConfiguration(ConnectionConfiguration connectionConfiguration) {
      checkArgument(connectionConfiguration != null, "connectionConfiguration can not be null");
      return builder().setConnectionConfiguration(connectionConfiguration).build();
    }

    /**
     * Provide a maximum size in number of documents for the batch see bulk API
     * (https:
     * docs (like Elasticsearch bulk size advice). See
     * https:
     * execution engine, size of bundles may vary, this sets the maximum size. Change this if you
     * need to have smaller Elasticsearch bulks.
     *
     * @param batchSize maximum batch size in number of documents
     * @return the {@link Write} with connection batch size set
     */
    public Write withMaxBatchSize(long batchSize) {
      checkArgument(batchSize > 0, "batchSize must be > 0, but was %s", batchSize);
      return builder().setMaxBatchSize(batchSize).build();
    }

    /**
     * Provide a maximum size in bytes for the batch see bulk API
     * (https:
     * (like Elasticsearch bulk size advice). See
     * https:
     * execution engine, size of bundles may vary, this sets the maximum size. Change this if you
     * need to have smaller Elasticsearch bulks.
     *
     * @param batchSizeBytes maximum batch size in bytes
     * @return the {@link Write} with connection batch size in bytes set
     */
    public Write withMaxBatchSizeBytes(long batchSizeBytes) {
      checkArgument(batchSizeBytes > 0, "batchSizeBytes must be > 0, but was %s", batchSizeBytes);
      return builder().setMaxBatchSizeBytes(batchSizeBytes).build();
    }

    /**
     * Provide a function to extract the id from the document. This id will be used as the document
     * id in Elasticsearch. Should the function throw an Exception then the batch will fail and the
     * exception propagated.
     *
     * @param idFn to extract the document ID
     * @return the {@link Write} with the function set
     */
    public Write withIdFn(FieldValueExtractFn idFn) {
      checkArgument(idFn != null, "idFn must not be null");
      return builder().setIdFn(idFn).build();
    }

    /**
     * Provide a function to extract the target index from the document allowing for dynamic
     * document routing. Should the function throw an Exception then the batch will fail and the
     * exception propagated.
     *
     * @param indexFn to extract the destination index from
     * @return the {@link Write} with the function set
     */
    public Write withIndexFn(FieldValueExtractFn indexFn) {
      checkArgument(indexFn != null, "indexFn must not be null");
      return builder().setIndexFn(indexFn).build();
    }

    /**
     * Provide a function to extract the target type from the document allowing for dynamic document
     * routing. Should the function throw an Exception then the batch will fail and the exception
     * propagated. Users are encouraged to consider carefully if multipe types are a sensible model
     * <a
     * href="https:
     * discussed in this blog</a>.
     *
     * @param typeFn to extract the destination index from
     * @return the {@link Write} with the function set
     */
    public Write withTypeFn(FieldValueExtractFn typeFn) {
      checkArgument(typeFn != null, "typeFn must not be null");
      return builder().setTypeFn(typeFn).build();
    }

    /**
     * Provide an instruction to control whether partial updates or inserts (default) are issued to
     * Elasticsearch.
     *
     * @param usePartialUpdate set to true to issue partial updates
     * @return the {@link Write} with the partial update control set
     */
    public Write withUsePartialUpdate(boolean usePartialUpdate) {
      return builder().setUsePartialUpdate(usePartialUpdate).build();
    }

    /**
     * Provides configuration to retry a failed batch call to Elasticsearch. A batch is considered
     * as failed if the underlying {@link RestClient} surfaces 429 HTTP status code as error for one
     * or more of the items in the {@link Response}. Users should consider that retrying might
     * compound the underlying problem which caused the initial failure. Users should also be aware
     * that once retrying is exhausted the error is surfaced to the runner which <em>may</em> then
     * opt to retry the current bundle in entirety or abort if the max number of retries of the
     * runner is completed. Retrying uses an exponential backoff algorithm, with minimum backoff of
     * 5 seconds and then surfacing the error once the maximum number of retries or maximum
     * configuration duration is exceeded.
     *
     * <p>Example use:
     *
     * <pre>{@code
     * ElasticsearchIO.write()
     *   .withRetryConfiguration(ElasticsearchIO.RetryConfiguration.create(10, Duration.standardMinutes(3))
     *   ...
     * }</pre>
     *
     * @param retryConfiguration the rules which govern the retry behavior
     * @return the {@link Write} with retrying configured
     */
    public Write withRetryConfiguration(RetryConfiguration retryConfiguration) {
      checkArgument(retryConfiguration != null, "retryConfiguration is required");
      return builder().setRetryConfiguration(retryConfiguration).build();
    }

    @Override
    public PDone expand(PCollection<String> input) {
      ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
      checkState(connectionConfiguration != null, "withConnectionConfiguration() is required");
      input.apply(ParDo.of(new WriteFn(this)));
      return PDone.in(input.getPipeline());
    }

    /** {@link DoFn} to for the {@link Write} transform. */
    @VisibleForTesting
    static class WriteFn extends DoFn<String, Void> {
      private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
      private static final int DEFAULT_RETRY_ON_CONFLICT = 5; 

      private static final Duration RETRY_INITIAL_BACKOFF = Duration.standardSeconds(5);

      @VisibleForTesting
      static final String RETRY_ATTEMPT_LOG = "Error writing to Elasticsearch. Retry attempt[%d]";

      @VisibleForTesting
      static final String RETRY_FAILED_LOG =
          "Error writing to ES after %d attempt(s). No more attempts allowed";

      private transient FluentBackoff retryBackoff;

      private int backendVersion;
      private final Write spec;
      private transient RestClient restClient;
      private ArrayList<String> batch;
      private long currentBatchSizeBytes;

      
      private static class DocumentMetadata implements Serializable {
        final String index;
        final String type;
        final String id;
        final Integer retryOnConflict;

        DocumentMetadata(String index, String type, String id, Integer retryOnConflict) {
          this.index = index;
          this.type = type;
          this.id = id;
          this.retryOnConflict = retryOnConflict;
        }
      }

      @VisibleForTesting
      WriteFn(Write spec) {
        this.spec = spec;
      }

      @Setup
      public void setup() throws IOException {
        ConnectionConfiguration connectionConfiguration = spec.getConnectionConfiguration();
        backendVersion = getBackendVersion(connectionConfiguration);
        restClient = connectionConfiguration.createClient();

        retryBackoff =
            FluentBackoff.DEFAULT.withMaxRetries(0).withInitialBackoff(RETRY_INITIAL_BACKOFF);

        if (spec.getRetryConfiguration() != null) {
          retryBackoff =
              FluentBackoff.DEFAULT
                  .withInitialBackoff(RETRY_INITIAL_BACKOFF)
                  .withMaxRetries(spec.getRetryConfiguration().getMaxAttempts() - 1)
                  .withMaxCumulativeBackoff(spec.getRetryConfiguration().getMaxDuration());
        }
      }

      @StartBundle
      public void startBundle(StartBundleContext context) {
        batch = new ArrayList<>();
        currentBatchSizeBytes = 0;
      }

      private class DocumentMetadataSerializer extends StdSerializer<DocumentMetadata> {

        private DocumentMetadataSerializer() {
          super(DocumentMetadata.class);
        }

        @Override
        public void serialize(
            DocumentMetadata value, JsonGenerator gen, SerializerProvider provider)
            throws IOException {
          gen.writeStartObject();
          if (value.index != null) {
            gen.writeStringField("_index", value.index);
          }
          if (value.type != null) {
            gen.writeStringField("_type", value.type);
          }
          if (value.id != null) {
            gen.writeStringField("_id", value.id);
          }
          if (value.retryOnConflict != null && (backendVersion <= 6)) {
            gen.writeNumberField("_retry_on_conflict", value.retryOnConflict);
          }
          if (value.retryOnConflict != null && backendVersion >= 7) {
            gen.writeNumberField("retry_on_conflict", value.retryOnConflict);
          }
          gen.writeEndObject();
        }
      }
      /**
       * Extracts the components that comprise the document address from the document using the
       * {@link FieldValueExtractFn} configured. This allows any or all of the index, type and
       * document id to be controlled on a per document basis. If none are provided then an empty
       * default of {@code {}} is returned. Sanitization of the index is performed, automatically
       * lower-casing the value as required by Elasticsearch.
       *
       * @param document the json from which the index, type and id may be extracted
       * @return the document address as JSON or the default
       * @throws IOException if the document cannot be parsed as JSON
       */
      private String getDocumentMetadata(String document) throws IOException {
        if (spec.getIndexFn() != null || spec.getTypeFn() != null || spec.getIdFn() != null) {
          
          JsonNode parsedDocument = OBJECT_MAPPER.readTree(document);

          DocumentMetadata metadata =
              new DocumentMetadata(
                  spec.getIndexFn() != null
                      ? lowerCaseOrNull(spec.getIndexFn().apply(parsedDocument))
                      : null,
                  spec.getTypeFn() != null ? spec.getTypeFn().apply(parsedDocument) : null,
                  spec.getIdFn() != null ? spec.getIdFn().apply(parsedDocument) : null,
                  spec.getUsePartialUpdate() ? DEFAULT_RETRY_ON_CONFLICT : null);
          SimpleModule module = new SimpleModule();
          module.addSerializer(DocumentMetadata.class, new DocumentMetadataSerializer());
          OBJECT_MAPPER.registerModule(module);
          return OBJECT_MAPPER.writeValueAsString(metadata);
        } else {
          return "{}"; 
        }
      }

      private static String lowerCaseOrNull(String input) {
        return input == null ? null : input.toLowerCase();
      }

      @ProcessElement
      public void processElement(ProcessContext context) throws Exception {
        String document = context.element();
        String documentMetadata = getDocumentMetadata(document);

        
        if (spec.getUsePartialUpdate()) {
          batch.add(
              String.format(
                  "{ \"update\" : %s }%n{ \"doc\" : %s, \"doc_as_upsert\" : true }%n",
                  documentMetadata, document));
        } else {
          batch.add(String.format("{ \"index\" : %s }%n%s%n", documentMetadata, document));
        }

        currentBatchSizeBytes += document.getBytes(StandardCharsets.UTF_8).length;
        if (batch.size() >= spec.getMaxBatchSize()
            || currentBatchSizeBytes >= spec.getMaxBatchSizeBytes()) {
          flushBatch();
        }
      }

      @FinishBundle
      public void finishBundle(FinishBundleContext context)
          throws IOException, InterruptedException {
        flushBatch();
      }

      private void flushBatch() throws IOException, InterruptedException {
        if (batch.isEmpty()) {
          return;
        }
        StringBuilder bulkRequest = new StringBuilder();
        for (String json : batch) {
          bulkRequest.append(json);
        }
        batch.clear();
        currentBatchSizeBytes = 0;
        Response response;
        HttpEntity responseEntity;
        
        
        
        String endPoint =
            String.format(
                "/%s/%s/_bulk",
                spec.getConnectionConfiguration().getIndex(),
                spec.getConnectionConfiguration().getType());
        HttpEntity requestBody =
            new NStringEntity(bulkRequest.toString(), ContentType.APPLICATION_JSON);
        Request request = new Request("POST", endPoint);
        request.addParameters(Collections.emptyMap());
        request.setEntity(requestBody);
        response = restClient.performRequest(request);
        responseEntity = new BufferedHttpEntity(response.getEntity());
        if (spec.getRetryConfiguration() != null
            && spec.getRetryConfiguration().getRetryPredicate().test(responseEntity)) {
          responseEntity = handleRetry("POST", endPoint, Collections.emptyMap(), requestBody);
        }
        checkForErrors(responseEntity, backendVersion, spec.getUsePartialUpdate());
      }

      /** retry request based on retry configuration policy. */
      private HttpEntity handleRetry(
          String method, String endpoint, Map<String, String> params, HttpEntity requestBody)
          throws IOException, InterruptedException {
        Response response;
        HttpEntity responseEntity;
        Sleeper sleeper = Sleeper.DEFAULT;
        BackOff backoff = retryBackoff.backoff();
        int attempt = 0;
        
        while (BackOffUtils.next(sleeper, backoff)) {
          LOG.warn(String.format(RETRY_ATTEMPT_LOG, ++attempt));
          Request request = new Request(method, endpoint);
          request.addParameters(params);
          request.setEntity(requestBody);
          response = restClient.performRequest(request);
          responseEntity = new BufferedHttpEntity(response.getEntity());
          
          if (!spec.getRetryConfiguration().getRetryPredicate().test(responseEntity)) {
            return responseEntity;
          }
        }
        throw new IOException(String.format(RETRY_FAILED_LOG, attempt));
      }

      @Teardown
      public void closeClient() throws IOException {
        if (restClient != null) {
          restClient.close();
        }
      }
    }
  }

  static int getBackendVersion(ConnectionConfiguration connectionConfiguration) {
    try (RestClient restClient = connectionConfiguration.createClient()) {
      Request request = new Request("GET", "");
      Response response = restClient.performRequest(request);
      JsonNode jsonNode = parseResponse(response.getEntity());
      int backendVersion =
          Integer.parseInt(jsonNode.path("version").path("number").asText().substring(0, 1));
      checkArgument(
          (backendVersion == 2
              || backendVersion == 5
              || backendVersion == 6
              || backendVersion == 7),
          "The Elasticsearch version to connect to is %s.x. "
              + "This version of the ElasticsearchIO is only compatible with "
              + "Elasticsearch v7.x, v6.x, v5.x and v2.x",
          backendVersion);
      return backendVersion;

    } catch (IOException e) {
      throw new IllegalArgumentException("Cannot get Elasticsearch version", e);
    }
  }
}