Merge pull request #21914 from [cherry-pick][release-2.40.0] Merge pull request #21895: Drops usage of setWindowingStrategyInterna…

[cherry-pick][release-2.40.0] Merge pull request #21895: Drops usage of setWindowingStrategyInterna…

Author: pabloem

sdks/java/io/elasticsearch/src/main/java/org/apache/beam/sdk/io/elasticsearch/ElasticsearchIO.java

@@ -75,6 +75,7 @@
 import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
 import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.transforms.windowing.WindowFn;
 import org.apache.beam.sdk.util.BackOff;
 import org.apache.beam.sdk.util.BackOffUtils;
 import org.apache.beam.sdk.util.FluentBackoff;
@@ -85,7 +86,6 @@
 import org.apache.beam.sdk.values.PCollectionTuple;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
-import org.apache.beam.sdk.values.WindowingStrategy;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings;
 import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Streams;
@@ -2279,10 +2279,13 @@ public PCollectionTuple expand(PCollection<Document> input) {
       ConnectionConfiguration connectionConfiguration = getConnectionConfiguration();
       checkState(connectionConfiguration != null, "withConnectionConfiguration() is required");
 
-      WindowingStrategy<?, ?> originalStrategy = input.getWindowingStrategy();
+      @SuppressWarnings("unchecked")
+      WindowFn<Document, ?> originalWindowFn =
+          (WindowFn<Document, ?>) input.getWindowingStrategy().getWindowFn();
 
       PCollection<Document> docResults;
-      PCollection<Document> globalDocs = input.apply(Window.into(new GlobalWindows()));
+      PCollection<Document> globalDocs =
+          input.apply("Window inputs globally", Window.into(new GlobalWindows()));
 
       if (getUseStatefulBatches()) {
         docResults =
@@ -2294,7 +2297,8 @@ public PCollectionTuple expand(PCollection<Document> input) {
       }
 
       return docResults
-          .setWindowingStrategyInternal(originalStrategy)
+          // Restore windowing of input
+          .apply("Restore original windows", Window.into(originalWindowFn))
           .apply(
               ParDo.of(new ResultFilteringFn())
                   .withOutputTags(Write.SUCCESSFUL_WRITES, TupleTagList.of(Write.FAILED_WRITES)));