From 10579a144cbd2aa8b861406c251c436ddd7e7919 Mon Sep 17 00:00:00 2001
From: Sam Whittle <samuelw@google.com>
Date: Mon, 2 Jun 2025 20:13:03 +0000
Subject: [PATCH] Explicitly handle singleton iterators instead of using helper
 and catching exceptions which may be from generating iterable

---
 .../beam/sdk/values/PCollectionViews.java     | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
index 6df68ea964d6..05612cb9fc77 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/values/PCollectionViews.java
@@ -540,14 +540,16 @@ public Materialization<IterableView<T>> getMaterialization() {
 
     @Override
     public T apply(IterableView<T> primitiveViewT) {
-      try {
-        return Iterables.getOnlyElement(primitiveViewT.get());
-      } catch (NoSuchElementException exc) {
+      Iterator<T> iterator = primitiveViewT.get().iterator();
+      if (!iterator.hasNext()) {
         return getDefaultValue();
-      } catch (IllegalArgumentException exc) {
+      }
+      T result = iterator.next();
+      if (iterator.hasNext()) {
         throw new IllegalArgumentException(
             "PCollection with more than one element accessed as a singleton view.");
       }
+      return result;
     }
 
     @Override
@@ -638,14 +640,16 @@ public Materialization<MultimapView<Void, T>> getMaterialization() {
 
     @Override
     public T apply(MultimapView<Void, T> primitiveViewT) {
-      try {
-        return Iterables.getOnlyElement(primitiveViewT.get(null));
-      } catch (NoSuchElementException exc) {
+      Iterator<T> iterator = primitiveViewT.get(null).iterator();
+      if (!iterator.hasNext()) {
         return getDefaultValue();
-      } catch (IllegalArgumentException exc) {
+      }
+      T result = iterator.next();
+      if (iterator.hasNext()) {
         throw new IllegalArgumentException(
             "PCollection with more than one element accessed as a singleton view.");
       }
+      return result;
     }
 
     @Override