You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@celeborn.apache.org by "Copilot (via GitHub)" <gi...@apache.org> on 2026/05/12 02:32:56 UTC
Re: [PR] [CELEBORN-2315] Add iterator fully-consumed validation after shuffle write [celeborn]
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: issues-unsubscribe@celeborn.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org