[Spark] Restore memory sensitive GBK translation (#33520)

runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java

@@ -146,32 +146,36 @@ public void evaluate(GroupByKey<K, V> transform, EvaluationContext context) {
 
         JavaRDD<WindowedValue<KV<K, Iterable<V>>>> groupedByKey;
         Partitioner partitioner = getPartitioner(context);
-        // As this is batch, we can ignore triggering and allowed lateness parameters.
-        if (windowingStrategy.getWindowFn().equals(new GlobalWindows())
-            && windowingStrategy.getTimestampCombiner().equals(TimestampCombiner.END_OF_WINDOW)) {
-          // we can drop the windows and recover them later
-          groupedByKey =
-              GroupNonMergingWindowsFunctions.groupByKeyInGlobalWindow(
-                  inRDD, keyCoder, coder.getValueCoder(), partitioner);
-        } else if (GroupNonMergingWindowsFunctions.isEligibleForGroupByWindow(windowingStrategy)) {
+        if (GroupNonMergingWindowsFunctions.isEligibleForGroupByWindow(windowingStrategy)) {
           // we can have a memory sensitive translation for non-merging windows
           groupedByKey =
               GroupNonMergingWindowsFunctions.groupByKeyAndWindow(
                   inRDD, keyCoder, coder.getValueCoder(), windowingStrategy, partitioner);
         } else {
-          // --- group by key only.
-          JavaRDD<KV<K, Iterable<WindowedValue<V>>>> groupedByKeyOnly =
-              GroupCombineFunctions.groupByKeyOnly(inRDD, keyCoder, wvCoder, partitioner);
 
-          // --- now group also by window.
-          // for batch, GroupAlsoByWindow uses an in-memory StateInternals.
-          groupedByKey =
-              groupedByKeyOnly.flatMap(
-                  new SparkGroupAlsoByWindowViaOutputBufferFn<>(
-                      windowingStrategy,
-                      new TranslationUtils.InMemoryStateInternalsFactory<>(),
-                      SystemReduceFn.buffering(coder.getValueCoder()),
-                      context.getSerializableOptions()));
+          // As this is batch, we can ignore triggering and allowed lateness parameters.
+          if (windowingStrategy.getWindowFn().equals(new GlobalWindows())
+              && windowingStrategy.getTimestampCombiner().equals(TimestampCombiner.END_OF_WINDOW)) {
+
+            // we can drop the windows and recover them later
+            groupedByKey =
+                GroupNonMergingWindowsFunctions.groupByKeyInGlobalWindow(
+                    inRDD, keyCoder, coder.getValueCoder(), partitioner);
+          } else {
+            // --- group by key only.
+            JavaRDD<KV<K, Iterable<WindowedValue<V>>>> groupedByKeyOnly =
+                GroupCombineFunctions.groupByKeyOnly(inRDD, keyCoder, wvCoder, partitioner);
+
+            // --- now group also by window.
+            // for batch, GroupAlsoByWindow uses an in-memory StateInternals.
+            groupedByKey =
+                groupedByKeyOnly.flatMap(
+                    new SparkGroupAlsoByWindowViaOutputBufferFn<>(
+                        windowingStrategy,
+                        new TranslationUtils.InMemoryStateInternalsFactory<>(),
+                        SystemReduceFn.buffering(coder.getValueCoder()),
+                        context.getSerializableOptions()));
+          }
         }
         context.putDataset(transform, new BoundedDataset<>(groupedByKey));
       }