Copilot commented on code in PR #3672:
URL: https://github.com/apache/celeborn/pull/3672#discussion_r3223347240


##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java:
##########
@@ -212,26 +212,33 @@ public long getPeakMemoryUsedBytes() {
     return peakMemoryUsedBytes;
   }
 
-  void doWrite(scala.collection.Iterator<Product2<K, V>> records) throws 
IOException {
+  // Returns true if the iterator still has records (i.e., not fully consumed)
+  @VisibleForTesting
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records) throws 
IOException {
     if (canUseFastWrite()) {
       fastWrite0(records);
+      return records.hasNext();
     } else if (dep.mapSideCombine()) {
       if (dep.aggregator().isEmpty()) {
         throw new UnsupportedOperationException(
             "When using map side combine, an aggregator must be specified.");
       }
-      write0(dep.aggregator().get().combineValuesByKey(records, taskContext));
+      scala.collection.Iterator combinedIterator =

Review Comment:
   `combinedIterator` is declared as a raw `scala.collection.Iterator`, which 
will introduce unchecked/rawtype warnings and hides the expected element type. 
Consider declaring it with a wildcarded type (e.g., 
`scala.collection.Iterator<Product2<K, ?>>`), since `write0` already treats 
values as `?`.
   



##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java:
##########
@@ -167,6 +157,25 @@ public void write(scala.collection.Iterator<Product2<K, 
V>> records) throws IOEx
     }
   }
 
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records) throws 
IOException {
+    if (canUseFastWrite()) {
+      fastWrite0(records);
+      return records.hasNext();
+    } else if (dep.mapSideCombine()) {
+      if (dep.aggregator().isEmpty()) {
+        throw new UnsupportedOperationException(
+            "When using map side combine, an aggregator must be specified.");
+      }
+      scala.collection.Iterator combinedIterator =
+          dep.aggregator().get().combineValuesByKey(records, taskContext);
+      write0(combinedIterator);
+      return combinedIterator.hasNext();
+    } else {
+      write0(records);
+      return records.hasNext();
+    }
+  }

Review Comment:
   `doWrite(...)` is only used internally by `write(...)` in this class; making 
it package-private increases API surface without a clear need. Consider making 
it `private`.



##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java:
##########
@@ -359,6 +369,8 @@ private void close() throws IOException, 
InterruptedException {
 
     updateMapStatus();
 
+    SparkUtils.assertIteratorFullyConsumed(iteratorHasNext);
+
     sendBufferPool.returnBuffer(sendBuffers);
     sendBuffers = null;
     sendOffsets = null;

Review Comment:
   `assertIteratorFullyConsumed(iteratorHasNext)` is called before returning 
`sendBuffers`/nulling fields. If the assertion triggers (TaskKilledException), 
buffers/arrays won’t be returned/cleared in this method. Consider returning 
buffers (and nulling state) before the assertion, or ensuring the failure path 
also releases these resources.



##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java:
##########
@@ -184,6 +174,26 @@ public void write(scala.collection.Iterator<Product2<K, 
V>> records) throws IOEx
     }
   }
 
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records)
+      throws IOException, InterruptedException {
+    if (canUseFastWrite()) {
+      fastWrite0(records);
+      return records.hasNext();
+    } else if (dep.mapSideCombine()) {
+      if (dep.aggregator().isEmpty()) {
+        throw new UnsupportedOperationException(
+            "When using map side combine, an aggregator must be specified.");
+      }
+      scala.collection.Iterator combinedIterator =

Review Comment:
   `combinedIterator` is declared as a raw `scala.collection.Iterator`, which 
can lead to unchecked/rawtype warnings. Prefer a parameterized type (e.g., 
`scala.collection.Iterator<Product2<K, ?>>`) to preserve type information while 
still allowing map-side-combine's value type to vary.
   



##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java:
##########
@@ -188,6 +178,26 @@ public void write(scala.collection.Iterator<Product2<K, 
V>> records) throws IOEx
     }
   }
 
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records)
+      throws IOException, InterruptedException {
+    if (canUseFastWrite()) {
+      fastWrite0(records);
+      return records.hasNext();
+    } else if (dep.mapSideCombine()) {
+      if (dep.aggregator().isEmpty()) {
+        throw new UnsupportedOperationException(
+            "When using map side combine, an aggregator must be specified.");
+      }
+      scala.collection.Iterator combinedIterator =
+          dep.aggregator().get().combineValuesByKey(records, taskContext);
+      write0(combinedIterator);
+      return combinedIterator.hasNext();
+    } else {

Review Comment:
   `combinedIterator` is declared as a raw `scala.collection.Iterator`, which 
can lead to unchecked/rawtype warnings. Prefer a parameterized/wildcard type 
(e.g., `scala.collection.Iterator<Product2<K, ?>>`) to keep typing explicit.



##########
client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java:
##########
@@ -184,6 +174,26 @@ public void write(scala.collection.Iterator<Product2<K, 
V>> records) throws IOEx
     }
   }
 
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records)

Review Comment:
   `doWrite(...)` is only called from within this class; keeping it 
package-private unnecessarily expands the surface area. Consider making it 
`private` (or annotating with `@VisibleForTesting` if it’s intended for tests).
   



##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java:
##########
@@ -167,6 +157,25 @@ public void write(scala.collection.Iterator<Product2<K, 
V>> records) throws IOEx
     }
   }
 
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records) throws 
IOException {
+    if (canUseFastWrite()) {
+      fastWrite0(records);
+      return records.hasNext();
+    } else if (dep.mapSideCombine()) {
+      if (dep.aggregator().isEmpty()) {
+        throw new UnsupportedOperationException(
+            "When using map side combine, an aggregator must be specified.");
+      }
+      scala.collection.Iterator combinedIterator =

Review Comment:
   `combinedIterator` is declared as a raw `scala.collection.Iterator`, which 
can produce unchecked/rawtype warnings. Consider using a parameterized/wildcard 
type (e.g., `scala.collection.Iterator<Product2<K, ?>>`) for clearer typing.
   



##########
client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java:
##########
@@ -188,6 +178,26 @@ public void write(scala.collection.Iterator<Product2<K, 
V>> records) throws IOEx
     }
   }
 
+  boolean doWrite(scala.collection.Iterator<Product2<K, V>> records)
+      throws IOException, InterruptedException {
+    if (canUseFastWrite()) {
+      fastWrite0(records);
+      return records.hasNext();
+    } else if (dep.mapSideCombine()) {
+      if (dep.aggregator().isEmpty()) {
+        throw new UnsupportedOperationException(
+            "When using map side combine, an aggregator must be specified.");
+      }
+      scala.collection.Iterator combinedIterator =
+          dep.aggregator().get().combineValuesByKey(records, taskContext);
+      write0(combinedIterator);
+      return combinedIterator.hasNext();
+    } else {
+      write0(records);
+      return records.hasNext();
+    }
+  }

Review Comment:
   `doWrite(...)` is package-private but only invoked from `write(...)` inside 
this class. Consider making it `private` to avoid exposing extra methods on a 
core shuffle writer implementation.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to