This is an automated email from the ASF dual-hosted git repository.
zaynt4606 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 47bf69ff4e [CELEBORN-2480][FOLLOWUP] Test columnar batch limits and
round trips
47bf69ff4e is described below
commit 47bf69ff4e0b96d6e34ee93d6b864b5800155942
Author: ForVic <[email protected]>
AuthorDate: Wed Oct 21 13:46:58 2026 +0800
[CELEBORN-2480][FOLLOWUP] Test columnar batch limits and round trips
### What changes were proposed in this pull request?
Add focused tests for byte-limited columnar shuffle batches in the Spark 3,
3.5 and 4 suites:
- The row limit wins when reached before the byte limit.
- Interleaved partitions keep independent byte counts and reset them after
each push.
- Pushed batches decode to the original rows, in order, with plain, codegen
and dictionary writers.
### Why are the changes needed?
Follow up on the [test coverage feedback on
#3873](https://github.com/apache/celeborn/pull/3873#issuecomment-6282765823).
The existing test checked batch row counts for one partition; the helper now
also checks decoded contents.
### Does this PR resolve a correctness bug?
- [ ] Yes
No known bug; regression coverage only.
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
No production code changes.
### How was this patch tested?
`ColumnarHashBasedShuffleWriterSuiteJ`: 15 tests passed on each of Spark
3.3, 3.5 and 4.0 (JDK 17). Spotless checks passed for all three modules.
```sh
build/mvn -P<profile> -pl client-spark/<module> -am test \
-Dtest=ColumnarHashBasedShuffleWriterSuiteJ -DwildcardSuites=none
```
Profiles/modules: `spark-3.3` / `spark-3-columnar-shuffle`, `spark-3.5` /
`spark-3.5-columnar-shuffle`, `spark-4.0` / `spark-4-columnar-shuffle`.
Closes #3887 from ForVic/columnar-batch-test-coverage.
Authored-by: ForVic <[email protected]>
Signed-off-by: zhengtao <[email protected]>
AI-Contributed/Feature: 0/0
AI-Contributed/UT: 0/243
---
.../ColumnarHashBasedShuffleWriterSuiteJ.java | 81 +++++++++++++++++++---
.../ColumnarHashBasedShuffleWriterSuiteJ.java | 81 +++++++++++++++++++---
.../ColumnarHashBasedShuffleWriterSuiteJ.java | 81 +++++++++++++++++++---
3 files changed, 219 insertions(+), 24 deletions(-)
diff --git
a/client-spark/spark-3-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
b/client-spark/spark-3-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
index 4cc24575c4..4b50211379 100644
---
a/client-spark/spark-3-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
+++
b/client-spark/spark-3-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
@@ -18,11 +18,18 @@
package org.apache.spark.shuffle.celeborn;
import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
+import java.io.ByteArrayInputStream;
import java.io.File;
import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
import java.util.UUID;
import scala.Product2;
@@ -34,6 +41,7 @@ import org.apache.spark.HashPartitioner;
import org.apache.spark.Partitioner;
import org.apache.spark.SparkVersionUtil;
import org.apache.spark.TaskContext;
+import org.apache.spark.serializer.DeserializationStream;
import org.apache.spark.serializer.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
@@ -110,9 +118,48 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
// Without it, only the row limit applies and all rows stay in one batch.
assertArrayEquals(new int[] {30}, writeAndGetBatchRowCounts(new
CelebornConf(), 30));
+
+ // The row limit still applies when byte-based batching is enabled.
+ conf.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE().key(), "3");
+ assertArrayEquals(new int[] {3, 3, 3, 1}, writeAndGetBatchRowCounts(conf,
10));
+ }
+
+ @Test
+ public void testInterleavedByteLimitedBatchesRoundTrip() throws Exception {
+ // Every third row goes to partition 1 (10 rows) and the rest to partition
0 (18 rows). The
+ // uneven, interleaved split exercises independent byte counts, resets and
final partial
+ // batches.
+ final int[] partitions = new int[28];
+ for (int i = 0; i < partitions.length; i++) {
+ partitions[i] = i % 3 == 0 ? 1 : 0;
+ }
+ for (String mode : new String[] {"plain", "codegen", "dictionary"}) {
+ final CelebornConf conf =
+ new CelebornConf()
+
.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_ENABLED().key(), "true")
+
.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_THRESHOLD().key(), "1k")
+ .set(
+ CelebornConf.COLUMNAR_SHUFFLE_CODEGEN_ENABLED().key(),
+ Boolean.toString(mode.equals("codegen")))
+ .set(
+
CelebornConf.COLUMNAR_SHUFFLE_DICTIONARY_ENCODING_ENABLED().key(),
+ Boolean.toString(mode.equals("dictionary")));
+ Map<Integer, List<byte[]>> batches = writeAndCheckBatches(conf,
partitions);
+ assertArrayEquals(mode, new int[] {8, 8, 2},
batchRowCounts(batches.get(0)));
+ assertArrayEquals(mode, new int[] {8, 2},
batchRowCounts(batches.get(1)));
+ }
+ }
+
+ private int[] batchRowCounts(List<byte[]> batches) {
+ return batches.stream().mapToInt(batch ->
ByteBuffer.wrap(batch).getInt()).toArray();
}
private int[] writeAndGetBatchRowCounts(CelebornConf conf, int numRows)
throws Exception {
+ return batchRowCounts(writeAndCheckBatches(conf, new int[numRows]).get(0));
+ }
+
+ private Map<Integer, List<byte[]>> writeAndCheckBatches(CelebornConf conf,
int[] partitions)
+ throws Exception {
final String partitionIdPassthroughClazz =
SparkVersionUtil.isGreaterThan(3, 3)
? "org.apache.spark.PartitionIdPassthrough"
@@ -138,20 +185,38 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
metrics.shuffleWriteMetrics());
final UnsafeProjection projection = UnsafeProjection.create(schema);
- final UTF8String value = UTF8String.fromBytes(new byte[100]);
+ final Map<Integer, List<UnsafeRow>> expected = new HashMap<>();
final ListBuffer<Product2<Integer, UnsafeRow>> records = new
ListBuffer<>();
- for (int i = 0; i < numRows; i++) {
- UnsafeRow row = projection.apply(new GenericInternalRow(new Object[] {i,
value}));
- records.$plus$eq(new Tuple2<>(0, row.copy()));
+ for (int i = 0; i < partitions.length; i++) {
+ // Repeated 100-byte strings keep rows at 128 bytes and allow dictionary
encoding.
+ UTF8String value = UTF8String.fromString(String.format(Locale.ROOT,
"%0100d", i % 2));
+ UnsafeRow row = projection.apply(new GenericInternalRow(new Object[] {i,
value})).copy();
+ records.$plus$eq(new Tuple2<>(partitions[i], row));
+ expected.computeIfAbsent(partitions[i], ignored -> new
ArrayList<>()).add(row);
}
writer.write((Iterator) records.toIterator());
writer.stop(true);
client.shutdown();
- // mapId=0 and attemptId=0 from the mock TaskContext. Each batch starts
with its row count.
- return client.getPushDataByPartition(Utils.makeMapKey(shuffleId, 0,
0)).get(0).stream()
- .mapToInt(batch -> ByteBuffer.wrap(batch).getInt())
- .toArray();
+ // mapId=0 and attemptId=0 from the mock TaskContext.
+ Map<Integer, List<byte[]>> batches =
+ client.getPushDataByPartition(Utils.makeMapKey(shuffleId, 0, 0));
+ assertEquals(expected.keySet(), batches.keySet());
+ final SerializerInstance deserializer = newSerializerInstance(new
UnsafeRowSerializer(2, null));
+ for (Map.Entry<Integer, List<byte[]>> partition : batches.entrySet()) {
+ List<UnsafeRow> actual = new ArrayList<>();
+ for (byte[] batch : partition.getValue()) {
+ try (DeserializationStream stream =
+ deserializer.deserializeStream(new ByteArrayInputStream(batch))) {
+ Iterator<Tuple2<Object, Object>> rows = stream.asKeyValueIterator();
+ while (rows.hasNext()) {
+ actual.add(((UnsafeRow) rows.next()._2()).copy());
+ }
+ }
+ }
+ assertEquals("partition " + partition.getKey(),
expected.get(partition.getKey()), actual);
+ }
+ return batches;
}
@Override
diff --git
a/client-spark/spark-3.5-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
b/client-spark/spark-3.5-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
index 4cc24575c4..4b50211379 100644
---
a/client-spark/spark-3.5-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
+++
b/client-spark/spark-3.5-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
@@ -18,11 +18,18 @@
package org.apache.spark.shuffle.celeborn;
import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
+import java.io.ByteArrayInputStream;
import java.io.File;
import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
import java.util.UUID;
import scala.Product2;
@@ -34,6 +41,7 @@ import org.apache.spark.HashPartitioner;
import org.apache.spark.Partitioner;
import org.apache.spark.SparkVersionUtil;
import org.apache.spark.TaskContext;
+import org.apache.spark.serializer.DeserializationStream;
import org.apache.spark.serializer.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
@@ -110,9 +118,48 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
// Without it, only the row limit applies and all rows stay in one batch.
assertArrayEquals(new int[] {30}, writeAndGetBatchRowCounts(new
CelebornConf(), 30));
+
+ // The row limit still applies when byte-based batching is enabled.
+ conf.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE().key(), "3");
+ assertArrayEquals(new int[] {3, 3, 3, 1}, writeAndGetBatchRowCounts(conf,
10));
+ }
+
+ @Test
+ public void testInterleavedByteLimitedBatchesRoundTrip() throws Exception {
+ // Every third row goes to partition 1 (10 rows) and the rest to partition
0 (18 rows). The
+ // uneven, interleaved split exercises independent byte counts, resets and
final partial
+ // batches.
+ final int[] partitions = new int[28];
+ for (int i = 0; i < partitions.length; i++) {
+ partitions[i] = i % 3 == 0 ? 1 : 0;
+ }
+ for (String mode : new String[] {"plain", "codegen", "dictionary"}) {
+ final CelebornConf conf =
+ new CelebornConf()
+
.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_ENABLED().key(), "true")
+
.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_THRESHOLD().key(), "1k")
+ .set(
+ CelebornConf.COLUMNAR_SHUFFLE_CODEGEN_ENABLED().key(),
+ Boolean.toString(mode.equals("codegen")))
+ .set(
+
CelebornConf.COLUMNAR_SHUFFLE_DICTIONARY_ENCODING_ENABLED().key(),
+ Boolean.toString(mode.equals("dictionary")));
+ Map<Integer, List<byte[]>> batches = writeAndCheckBatches(conf,
partitions);
+ assertArrayEquals(mode, new int[] {8, 8, 2},
batchRowCounts(batches.get(0)));
+ assertArrayEquals(mode, new int[] {8, 2},
batchRowCounts(batches.get(1)));
+ }
+ }
+
+ private int[] batchRowCounts(List<byte[]> batches) {
+ return batches.stream().mapToInt(batch ->
ByteBuffer.wrap(batch).getInt()).toArray();
}
private int[] writeAndGetBatchRowCounts(CelebornConf conf, int numRows)
throws Exception {
+ return batchRowCounts(writeAndCheckBatches(conf, new int[numRows]).get(0));
+ }
+
+ private Map<Integer, List<byte[]>> writeAndCheckBatches(CelebornConf conf,
int[] partitions)
+ throws Exception {
final String partitionIdPassthroughClazz =
SparkVersionUtil.isGreaterThan(3, 3)
? "org.apache.spark.PartitionIdPassthrough"
@@ -138,20 +185,38 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
metrics.shuffleWriteMetrics());
final UnsafeProjection projection = UnsafeProjection.create(schema);
- final UTF8String value = UTF8String.fromBytes(new byte[100]);
+ final Map<Integer, List<UnsafeRow>> expected = new HashMap<>();
final ListBuffer<Product2<Integer, UnsafeRow>> records = new
ListBuffer<>();
- for (int i = 0; i < numRows; i++) {
- UnsafeRow row = projection.apply(new GenericInternalRow(new Object[] {i,
value}));
- records.$plus$eq(new Tuple2<>(0, row.copy()));
+ for (int i = 0; i < partitions.length; i++) {
+ // Repeated 100-byte strings keep rows at 128 bytes and allow dictionary
encoding.
+ UTF8String value = UTF8String.fromString(String.format(Locale.ROOT,
"%0100d", i % 2));
+ UnsafeRow row = projection.apply(new GenericInternalRow(new Object[] {i,
value})).copy();
+ records.$plus$eq(new Tuple2<>(partitions[i], row));
+ expected.computeIfAbsent(partitions[i], ignored -> new
ArrayList<>()).add(row);
}
writer.write((Iterator) records.toIterator());
writer.stop(true);
client.shutdown();
- // mapId=0 and attemptId=0 from the mock TaskContext. Each batch starts
with its row count.
- return client.getPushDataByPartition(Utils.makeMapKey(shuffleId, 0,
0)).get(0).stream()
- .mapToInt(batch -> ByteBuffer.wrap(batch).getInt())
- .toArray();
+ // mapId=0 and attemptId=0 from the mock TaskContext.
+ Map<Integer, List<byte[]>> batches =
+ client.getPushDataByPartition(Utils.makeMapKey(shuffleId, 0, 0));
+ assertEquals(expected.keySet(), batches.keySet());
+ final SerializerInstance deserializer = newSerializerInstance(new
UnsafeRowSerializer(2, null));
+ for (Map.Entry<Integer, List<byte[]>> partition : batches.entrySet()) {
+ List<UnsafeRow> actual = new ArrayList<>();
+ for (byte[] batch : partition.getValue()) {
+ try (DeserializationStream stream =
+ deserializer.deserializeStream(new ByteArrayInputStream(batch))) {
+ Iterator<Tuple2<Object, Object>> rows = stream.asKeyValueIterator();
+ while (rows.hasNext()) {
+ actual.add(((UnsafeRow) rows.next()._2()).copy());
+ }
+ }
+ }
+ assertEquals("partition " + partition.getKey(),
expected.get(partition.getKey()), actual);
+ }
+ return batches;
}
@Override
diff --git
a/client-spark/spark-4-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
b/client-spark/spark-4-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
index 4cc24575c4..4b50211379 100644
---
a/client-spark/spark-4-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
+++
b/client-spark/spark-4-columnar-shuffle/src/test/java/org/apache/spark/shuffle/celeborn/ColumnarHashBasedShuffleWriterSuiteJ.java
@@ -18,11 +18,18 @@
package org.apache.spark.shuffle.celeborn;
import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
+import java.io.ByteArrayInputStream;
import java.io.File;
import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
import java.util.UUID;
import scala.Product2;
@@ -34,6 +41,7 @@ import org.apache.spark.HashPartitioner;
import org.apache.spark.Partitioner;
import org.apache.spark.SparkVersionUtil;
import org.apache.spark.TaskContext;
+import org.apache.spark.serializer.DeserializationStream;
import org.apache.spark.serializer.KryoSerializer;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
@@ -110,9 +118,48 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
// Without it, only the row limit applies and all rows stay in one batch.
assertArrayEquals(new int[] {30}, writeAndGetBatchRowCounts(new
CelebornConf(), 30));
+
+ // The row limit still applies when byte-based batching is enabled.
+ conf.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE().key(), "3");
+ assertArrayEquals(new int[] {3, 3, 3, 1}, writeAndGetBatchRowCounts(conf,
10));
+ }
+
+ @Test
+ public void testInterleavedByteLimitedBatchesRoundTrip() throws Exception {
+ // Every third row goes to partition 1 (10 rows) and the rest to partition
0 (18 rows). The
+ // uneven, interleaved split exercises independent byte counts, resets and
final partial
+ // batches.
+ final int[] partitions = new int[28];
+ for (int i = 0; i < partitions.length; i++) {
+ partitions[i] = i % 3 == 0 ? 1 : 0;
+ }
+ for (String mode : new String[] {"plain", "codegen", "dictionary"}) {
+ final CelebornConf conf =
+ new CelebornConf()
+
.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_ENABLED().key(), "true")
+
.set(CelebornConf.COLUMNAR_SHUFFLE_BATCH_SIZE_IN_BYTES_THRESHOLD().key(), "1k")
+ .set(
+ CelebornConf.COLUMNAR_SHUFFLE_CODEGEN_ENABLED().key(),
+ Boolean.toString(mode.equals("codegen")))
+ .set(
+
CelebornConf.COLUMNAR_SHUFFLE_DICTIONARY_ENCODING_ENABLED().key(),
+ Boolean.toString(mode.equals("dictionary")));
+ Map<Integer, List<byte[]>> batches = writeAndCheckBatches(conf,
partitions);
+ assertArrayEquals(mode, new int[] {8, 8, 2},
batchRowCounts(batches.get(0)));
+ assertArrayEquals(mode, new int[] {8, 2},
batchRowCounts(batches.get(1)));
+ }
+ }
+
+ private int[] batchRowCounts(List<byte[]> batches) {
+ return batches.stream().mapToInt(batch ->
ByteBuffer.wrap(batch).getInt()).toArray();
}
private int[] writeAndGetBatchRowCounts(CelebornConf conf, int numRows)
throws Exception {
+ return batchRowCounts(writeAndCheckBatches(conf, new int[numRows]).get(0));
+ }
+
+ private Map<Integer, List<byte[]>> writeAndCheckBatches(CelebornConf conf,
int[] partitions)
+ throws Exception {
final String partitionIdPassthroughClazz =
SparkVersionUtil.isGreaterThan(3, 3)
? "org.apache.spark.PartitionIdPassthrough"
@@ -138,20 +185,38 @@ public class ColumnarHashBasedShuffleWriterSuiteJ extends
CelebornShuffleWriterS
metrics.shuffleWriteMetrics());
final UnsafeProjection projection = UnsafeProjection.create(schema);
- final UTF8String value = UTF8String.fromBytes(new byte[100]);
+ final Map<Integer, List<UnsafeRow>> expected = new HashMap<>();
final ListBuffer<Product2<Integer, UnsafeRow>> records = new
ListBuffer<>();
- for (int i = 0; i < numRows; i++) {
- UnsafeRow row = projection.apply(new GenericInternalRow(new Object[] {i,
value}));
- records.$plus$eq(new Tuple2<>(0, row.copy()));
+ for (int i = 0; i < partitions.length; i++) {
+ // Repeated 100-byte strings keep rows at 128 bytes and allow dictionary
encoding.
+ UTF8String value = UTF8String.fromString(String.format(Locale.ROOT,
"%0100d", i % 2));
+ UnsafeRow row = projection.apply(new GenericInternalRow(new Object[] {i,
value})).copy();
+ records.$plus$eq(new Tuple2<>(partitions[i], row));
+ expected.computeIfAbsent(partitions[i], ignored -> new
ArrayList<>()).add(row);
}
writer.write((Iterator) records.toIterator());
writer.stop(true);
client.shutdown();
- // mapId=0 and attemptId=0 from the mock TaskContext. Each batch starts
with its row count.
- return client.getPushDataByPartition(Utils.makeMapKey(shuffleId, 0,
0)).get(0).stream()
- .mapToInt(batch -> ByteBuffer.wrap(batch).getInt())
- .toArray();
+ // mapId=0 and attemptId=0 from the mock TaskContext.
+ Map<Integer, List<byte[]>> batches =
+ client.getPushDataByPartition(Utils.makeMapKey(shuffleId, 0, 0));
+ assertEquals(expected.keySet(), batches.keySet());
+ final SerializerInstance deserializer = newSerializerInstance(new
UnsafeRowSerializer(2, null));
+ for (Map.Entry<Integer, List<byte[]>> partition : batches.entrySet()) {
+ List<UnsafeRow> actual = new ArrayList<>();
+ for (byte[] batch : partition.getValue()) {
+ try (DeserializationStream stream =
+ deserializer.deserializeStream(new ByteArrayInputStream(batch))) {
+ Iterator<Tuple2<Object, Object>> rows = stream.asKeyValueIterator();
+ while (rows.hasNext()) {
+ actual.add(((UnsafeRow) rows.next()._2()).copy());
+ }
+ }
+ }
+ assertEquals("partition " + partition.getKey(),
expected.get(partition.getKey()), actual);
+ }
+ return batches;
}
@Override
