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]